openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-05 15:18:41 +00:00
Code Issues Packages Projects Releases Wiki Activity

code-mode: make callback shutdown locally bounded

Browse Source
This commit is contained in:
Channing Conger
2026-06-19 07:48:48 +00:00
parent a736ac14ad
commit f1d0826a54
5 changed files with 97 additions and 209 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
codex-rs/code-mode/src/cell_actor/callbacks.rs
View File

@@ -22,8 +22,14 @@ pub(super) fn spawn_notification<H: CellHost>(
cancellation_token: CancellationToken,
) {
tasks.spawn(async move {
if let Err(err) = host.notify(call_id, text, cancellation_token).await {
warn!("failed to deliver code mode notification: {err}");
tokio::select! {
biased;
_ = cancellation_token.cancelled() => {}
result = host.notify(call_id, text, cancellation_token.clone()) => {
if let Err(err) = result {
warn!("failed to deliver code mode notification: {err}");
}
}
}
});
}
@@ -37,9 +43,15 @@ pub(super) fn spawn_tool<H: CellHost>(
) {
tasks.spawn(async move {
let id = invocation.id.clone();
let command = match host.invoke_tool(invocation, cancellation_token).await {
Ok(result) => RuntimeCommand::ToolResponse { id, result },
Err(error_text) => RuntimeCommand::ToolError { id, error_text },
let command = tokio::select! {
biased;
_ = cancellation_token.cancelled() => {
return;
}
result = host.invoke_tool(invocation, cancellation_token.clone()) => match result {
Ok(result) => RuntimeCommand::ToolResponse { id, result },
Err(error_text) => RuntimeCommand::ToolError { id, error_text },
}
};
let _ = runtime_tx.send(command);
});

6
codex-rs/code-mode/src/cell_actor/types.rs
View File

@@ -34,8 +34,10 @@ pub(crate) struct CellToolCall {
/// Connects a cell actor to session-owned callbacks and lifecycle state.
///
/// Implementations must honor callback cancellation and must not return from
/// `closed` until the session can no longer route requests to the cell.
/// Implementations should forward callback cancellation to downstream work.
/// The actor stops awaiting callbacks once cancellation begins. Implementations
/// must not return from `closed` until the session can no longer route requests
/// to the cell.
pub(crate) trait CellHost: Send + Sync + 'static {
fn invoke_tool(
&self,

261
codex-rs/code-mode/src/service_contract_tests.rs
View File

@@ -1,6 +1,4 @@
use std::sync::Arc;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_protocol::ToolName;
@@ -17,22 +15,17 @@ use crate::ToolDefinition;
enum DelegateEvent {
NotificationStarted,
NotificationFinished,
NotificationCancelled,
ToolStarted,
ToolCancelled,
CellClosed(CellId),
}
struct BlockingDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
notification_finished: AtomicBool,
tool_finished: AtomicBool,
tool_release: Notify,
}
struct HeldNotificationDelegate {
struct NeverResolvingNotificationDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
notification_release: Notify,
}
struct ReleasableNotificationDelegate {
@@ -40,20 +33,14 @@ struct ReleasableNotificationDelegate {
notification_release: Notify,
}
impl HeldNotificationDelegate {
struct NeverResolvingToolDelegate {
events_tx: mpsc::UnboundedSender<DelegateEvent>,
}
impl NeverResolvingNotificationDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(
Arc::new(Self {
events_tx,
notification_release: Notify::new(),
}),
events_rx,
)
}
fn release_notification(&self) {
self.notification_release.notify_one();
(Arc::new(Self { events_tx }), events_rx)
}
}
@@ -74,18 +61,20 @@ impl ReleasableNotificationDelegate {
}
}
impl CodeModeSessionDelegate for HeldNotificationDelegate {
impl NeverResolvingToolDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(Arc::new(Self { events_tx }), events_rx)
}
}
impl CodeModeSessionDelegate for NeverResolvingNotificationDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
cancellation_token: CancellationToken,
_cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::ToolStarted);
cancellation_token.cancelled().await;
let _ = self.events_tx.send(DelegateEvent::ToolCancelled);
Err("cancelled".to_string())
})
Box::pin(async { Err("unexpected tool call".to_string()) })
}
fn notify<'a>(
@@ -93,14 +82,11 @@ impl CodeModeSessionDelegate for HeldNotificationDelegate {
_call_id: String,
_cell_id: CellId,
_text: String,
cancellation_token: CancellationToken,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
cancellation_token.cancelled().await;
let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);
self.notification_release.notified().await;
Ok(())
std::future::pending().await
})
}
@@ -138,7 +124,6 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
Ok(())
}
_ = cancellation_token.cancelled() => {
let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);
Err("cancelled".to_string())
}
}
@@ -152,14 +137,41 @@ impl CodeModeSessionDelegate for ReleasableNotificationDelegate {
}
}
impl CodeModeSessionDelegate for NeverResolvingToolDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
_cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::ToolStarted);
std::future::pending().await
})
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
_cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
Box::pin(async { Ok(()) })
}
fn cell_closed(&self, cell_id: &CellId) {
let _ = self
.events_tx
.send(DelegateEvent::CellClosed(cell_id.clone()));
}
}
impl BlockingDelegate {
fn new() -> (Arc<Self>, mpsc::UnboundedReceiver<DelegateEvent>) {
let (events_tx, events_rx) = mpsc::unbounded_channel();
(
Arc::new(Self {
events_tx,
notification_finished: AtomicBool::new(false),
tool_finished: AtomicBool::new(false),
tool_release: Notify::new(),
}),
events_rx,
@@ -181,12 +193,9 @@ impl CodeModeSessionDelegate for BlockingDelegate {
let _ = self.events_tx.send(DelegateEvent::ToolStarted);
tokio::select! {
_ = self.tool_release.notified() => {
self.tool_finished.store(true, Ordering::Release);
Ok(serde_json::Value::Null)
}
_ = cancellation_token.cancelled() => {
self.tool_finished.store(true, Ordering::Release);
let _ = self.events_tx.send(DelegateEvent::ToolCancelled);
Err("cancelled".to_string())
}
}
@@ -203,8 +212,6 @@ impl CodeModeSessionDelegate for BlockingDelegate {
Box::pin(async move {
let _ = self.events_tx.send(DelegateEvent::NotificationStarted);
cancellation_token.cancelled().await;
self.notification_finished.store(true, Ordering::Release);
let _ = self.events_tx.send(DelegateEvent::NotificationCancelled);
Err("cancelled".to_string())
})
}
@@ -389,7 +396,7 @@ async fn observed_natural_completion_wins_over_termination() {
}
#[tokio::test]
async fn termination_cancels_pending_callbacks_before_responding() {
async fn termination_discards_pending_callbacks_before_responding() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell = service
@@ -417,11 +424,6 @@ async fn termination_cancels_pending_callbacks_before_responding() {
content_items: Vec::new(),
})
);
assert!(delegate.notification_finished.load(Ordering::Acquire));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationCancelled
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
@@ -429,9 +431,9 @@ async fn termination_cancels_pending_callbacks_before_responding() {
}
#[tokio::test]
async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
let (delegate, mut events_rx) = HeldNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
async fn shutdown_does_not_await_notifications_during_natural_completion() {
let (delegate, mut events_rx) = NeverResolvingNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate));
service
.execute(execute_request(r#"notify("pending");"#))
.await
@@ -443,15 +445,12 @@ async fn shutdown_cancels_notifications_while_natural_completion_is_draining() {
);
let shutdown_service = Arc::clone(&service);
let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await });
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationCancelled
);
delegate.release_notification();
assert_eq!(shutdown.await.unwrap(), Ok(()));
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), async move {
shutdown_service.shutdown().await
})
.await
.expect("shutdown should not await a non-cooperative notification")
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
@@ -499,7 +498,7 @@ async fn natural_completion_waits_for_notifications_before_responding() {
}
#[tokio::test]
async fn termination_cancels_pending_tools_before_responding() {
async fn termination_discards_pending_tools_before_responding() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell = service
@@ -527,11 +526,6 @@ async fn termination_cancels_pending_tools_before_responding() {
content_items: Vec::new(),
})
);
assert!(delegate.tool_finished.load(Ordering::Acquire));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::ToolCancelled
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
@@ -539,7 +533,7 @@ async fn termination_cancels_pending_tools_before_responding() {
}
#[tokio::test]
async fn shutdown_cancels_pending_tools_before_returning() {
async fn shutdown_discards_pending_tools_before_returning() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
service
@@ -554,13 +548,7 @@ async fn shutdown_cancels_pending_tools_before_returning() {
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
let shutdown_service = Arc::clone(&service);
let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await });
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::ToolCancelled
);
assert_eq!(shutdown.await.unwrap(), Ok(()));
assert_eq!(shutdown_service.shutdown().await, Ok(()));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
@@ -568,98 +556,27 @@ async fn shutdown_cancels_pending_tools_before_returning() {
}
#[tokio::test]
async fn shutdown_cancels_mixed_callbacks_while_natural_completion_is_draining() {
let (delegate, mut events_rx) = HeldNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
async fn shutdown_completes_when_a_nested_tool_ignores_cancellation() {
let (delegate, mut events_rx) = NeverResolvingToolDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate));
service
.execute(ExecuteRequest {
enabled_tools: vec![blocking_tool()],
source: r#"notify("pending"); tools.block({});"#.to_string(),
source: r#"await tools.block({});"#.to_string(),
..execute_request("")
})
.await
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationStarted
);
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
let shutdown_service = Arc::clone(&service);
let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await });
let first_cancelled = next_event(&mut events_rx).await;
let second_cancelled = next_event(&mut events_rx).await;
assert!(matches!(
(&first_cancelled, &second_cancelled),
(
&DelegateEvent::NotificationCancelled,
&DelegateEvent::ToolCancelled,
) | (
&DelegateEvent::ToolCancelled,
&DelegateEvent::NotificationCancelled,
)
));
delegate.release_notification();
assert_eq!(shutdown.await.unwrap(), Ok(()));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}
#[tokio::test]
async fn repeated_termination_is_rejected_while_callback_cleanup_is_pending() {
let (delegate, mut events_rx) = HeldNotificationDelegate::new();
let service = Arc::new(CodeModeService::with_delegate(delegate.clone()));
let cell = service
.execute(execute_request(
r#"notify("pending"); await new Promise(() => {});"#,
))
.await
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationStarted
);
assert_eq!(
cell.initial_response().await.unwrap(),
RuntimeResponse::Yielded {
cell_id: cell_id("1"),
content_items: Vec::new(),
}
);
let terminating_service = Arc::clone(&service);
let first_termination =
tokio::spawn(async move { terminating_service.terminate(cell_id("1")).await });
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::NotificationCancelled
);
let shutdown_service = Arc::clone(&service);
let shutdown = tokio::spawn(async move { shutdown_service.shutdown().await });
let repeated_termination = service.terminate(cell_id("1")).await;
delegate.release_notification();
assert_eq!(
repeated_termination.unwrap_err(),
"exec cell 1 is already terminating"
);
assert_eq!(
first_termination.await.unwrap().unwrap(),
WaitOutcome::LiveCell(RuntimeResponse::Terminated {
cell_id: cell_id("1"),
content_items: Vec::new(),
})
);
assert_eq!(shutdown.await.unwrap(), Ok(()));
tokio::time::timeout(Duration::from_millis(/*millis*/ 100), async move {
shutdown_service.shutdown().await
})
.await
.expect("shutdown should not await a non-cooperative nested tool")
.unwrap();
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
@@ -760,39 +677,3 @@ async fn dropping_a_wait_allows_a_later_wait_to_observe_the_cell() {
})
);
}
#[tokio::test]
async fn natural_completion_cleans_up_callbacks_before_responding() {
let (delegate, mut events_rx) = BlockingDelegate::new();
let service = CodeModeService::with_delegate(delegate.clone());
let cell = service
.execute(ExecuteRequest {
enabled_tools: vec![blocking_tool()],
source: r#"tools.block({}); text("done");"#.to_string(),
yield_time_ms: Some(60_000),
..execute_request("")
})
.await
.unwrap();
assert_eq!(next_event(&mut events_rx).await, DelegateEvent::ToolStarted);
assert_eq!(
cell.initial_response().await.unwrap(),
RuntimeResponse::Result {
cell_id: cell_id("1"),
content_items: vec![FunctionCallOutputContentItem::InputText {
text: "done".to_string(),
}],
error_text: None,
}
);
assert!(delegate.tool_finished.load(Ordering::Acquire));
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::ToolCancelled
);
assert_eq!(
next_event(&mut events_rx).await,
DelegateEvent::CellClosed(cell_id("1"))
);
}

11
codex-rs/code-mode/src/session_runtime/tests.rs
View File

@@ -20,7 +20,6 @@ enum FirstCloseOutcome {
#[derive(Debug, PartialEq)]
enum NotificationEvent {
Started,
Cancelled,
Closed(CellId),
}
@@ -113,9 +112,7 @@ impl SessionRuntimeDelegate for BlockingNotificationDelegate {
.send(NotificationEvent::Started)
.map_err(|_| "test did not receive notification start".to_string())?;
cancellation_token.cancelled().await;
self.events_tx
.send(NotificationEvent::Cancelled)
.map_err(|_| "test did not receive notification cancellation".to_string())
Ok(())
}
async fn cell_closed(&self, cell_id: &CellId) -> Result<(), String> {
@@ -506,12 +503,6 @@ async fn drop_cancels_notifications_during_natural_completion_when_registry_is_l
drop(runtime);
drop(cells);
assert_eq!(
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), events_rx.recv())
.await
.unwrap(),
Some(NotificationEvent::Cancelled)
);
assert_eq!(
tokio::time::timeout(Duration::from_secs(/*secs*/ 1), events_rx.recv())
.await

6
codex-rs/code-mode/src/session_runtime/types.rs
View File

@@ -112,8 +112,10 @@ pub struct NestedToolCall {
/// Host callbacks used by cells owned by a [`super::SessionRuntime`].
///
/// Implementations must honor cancellation tokens and must not return from
/// `cell_closed` until downstream routing can no longer target the cell.
/// Implementations should forward callback cancellation tokens to downstream
/// work. The runtime stops awaiting callbacks once cancellation begins.
/// Implementations must not return from `cell_closed` until downstream routing
/// can no longer target the cell.
pub trait SessionRuntimeDelegate: Send + Sync + 'static {
fn invoke_tool(
&self,