openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-10-20 17:38:45 +00:00
Code Issues Packages Projects Releases Wiki Activity

Close remote delegate lifecycle races

Browse Source
This commit is contained in:
Channing Conger
2026-06-25 01:06:21 +00:00
parent 1add55e6a9
commit 1c685fb3ca
9 changed files with 319 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode-host/src/delegate.rs
+4
View File
@@ -12,6 +12,10 @@ use tokio_util::sync::CancellationToken;
use crate::peer::HostPeer;
/// Forwards delegate calls from one hosted session to the connected client.
///
/// The session ID associates each reverse request with the client-side delegate
/// that owns the corresponding logical session.
pub(super) struct RemoteDelegate {
session_id: SessionId,
peer: Arc<HostPeer>,
codex-rs/code-mode-host/src/lib.rs
+16 -3
View File
@@ -222,6 +222,10 @@ where
Ok(true)
}
/// Owns the sessions and in-flight operations for one client connection.
///
/// Disconnecting this state shuts down every hosted session before waiting for
/// its request tasks to finish.
struct HostState {
sessions: Mutex<HashMap<SessionId, Arc<InProcessCodeModeSession>>>,
seen_session_ids: Mutex<HashSet<SessionId>>,
@@ -268,13 +272,22 @@ impl HostState {
return;
}
};
let result = match self.session(&session_id) {
Ok(session) => session.execute_with_cell_id(cell_id.into(), request).await,
let session = match self.session(&session_id) {
Ok(session) => session,
Err(err) => {
self.respond(request_id, Err(err));
return;
}
};
let registration = match self.peer.register_cell(session_id, cell_id.clone().into())
{
Ok(registration) => registration,
Err(err) => {
self.respond(request_id, Err(err));
return;
}
};
let result = session.execute_with_cell_id(cell_id.into(), request).await;
match result {
Ok(started) => {
let cell_id = started.cell_id.clone();
@@ -284,7 +297,7 @@ impl HostState {
cell_id: cell_id.into(),
}),
);
self.peer.start_cell(session_id, request_id, started);
self.peer.start_cell(registration, request_id, started);
}
Err(err) => self.respond(request_id, Err(err)),
}
codex-rs/code-mode-host/src/peer.rs
+107 -59
View File
@@ -19,19 +19,30 @@ use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio_util::sync::CancellationToken;
/// Multiplexes host-to-client traffic for all sessions on one connection.
///
/// It correlates reverse delegate calls with their responses and routes each
/// cell's initial response, delegate requests, and closure notification through
/// a single cell driver.
pub(super) struct HostPeer {
outgoing_tx: mpsc::UnboundedSender<HostToClient>,
pending: Mutex<HashMap<DelegateRequestId, oneshot::Sender<Result<DelegateResponse, String>>>>,
next_request_id: AtomicI64,
disconnected: CancellationToken,
cell_routes: StdMutex<HashMap<(SessionId, CellId), CellRoute>>,
cell_routes: StdMutex<HashMap<(SessionId, CellId), mpsc::UnboundedSender<CellMessage>>>,
}
enum CellRoute {
Pending(Vec<CellMessage>),
Active(mpsc::UnboundedSender<CellMessage>),
/// Registers a cell's message channel before its runtime task can emit callbacks.
///
/// Dropping the registration removes the route, including when runtime admission
/// fails before the cell driver starts.
pub(super) struct CellRegistration {
peer: Arc<HostPeer>,
key: (SessionId, CellId),
messages_rx: mpsc::UnboundedReceiver<CellMessage>,
}
/// An event emitted by a hosted session for its per-cell driver to serialize.
enum CellMessage {
Delegate {
id: DelegateRequestId,
@@ -72,10 +83,17 @@ impl HostPeer {
DelegateRequest::InvokeTool { invocation } => invocation.cell_id.clone(),
DelegateRequest::Notify { cell_id, .. } => cell_id.clone(),
};
self.route_cell_message(
(session_id.clone(), cell_id.into()),
if !self.route_cell_message(
(session_id.clone(), cell_id.clone().into()),
CellMessage::Delegate { id, request },
);
) {
self.pending.lock().await.remove(&id);
pending.disarm();
return Err(format!(
"code-mode cell {} is not registered for session {session_id}",
cell_id.as_str()
));
}
tokio::select! {
response = response_rx => {
@@ -85,9 +103,11 @@ impl HostPeer {
})?
}
_ = cancellation_token.cancelled() => {
if self.pending.lock().await.remove(&id).is_some() {
let mut pending_calls = self.pending.lock().await;
if pending_calls.remove(&id).is_some() {
self.send(HostToClient::CancelDelegateRequest { id });
}
drop(pending_calls);
pending.disarm();
Err("code mode delegate request cancelled".to_string())
}
@@ -113,53 +133,64 @@ impl HostPeer {
self.disconnected.cancel();
}
pub(super) fn start_cell(
pub(super) fn register_cell(
self: &Arc<Self>,
session_id: SessionId,
request_id: RequestId,
started: StartedCell,
) {
let key = (session_id, started.cell_id.clone());
let (messages_tx, messages_rx) = mpsc::unbounded_channel();
let pending = self
.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner)
.insert(key.clone(), CellRoute::Active(messages_tx.clone()));
if let Some(CellRoute::Pending(messages)) = pending {
for message in messages {
let _ = messages_tx.send(message);
}
}
let peer = Arc::clone(self);
tokio::spawn(async move {
drive_cell(peer, key, request_id, started, messages_rx).await;
});
}
pub(super) fn close_cell(&self, session_id: SessionId, cell_id: CellId) {
self.route_cell_message((session_id, cell_id), CellMessage::Closed);
}
fn route_cell_message(&self, key: (SessionId, CellId), message: CellMessage) {
cell_id: CellId,
) -> Result<CellRegistration, String> {
use std::collections::hash_map::Entry;
let key = (session_id, cell_id);
let (messages_tx, messages_rx) = mpsc::unbounded_channel();
match self
.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner)
.entry(key)
.entry(key.clone())
{
Entry::Occupied(mut entry) => match entry.get_mut() {
CellRoute::Pending(messages) => messages.push(message),
CellRoute::Active(sender) => {
let _ = sender.send(message);
}
},
Entry::Occupied(_) => {
return Err(format!(
"code-mode cell {} is already registered for session {}",
key.1, key.0
));
}
Entry::Vacant(entry) => {
entry.insert(CellRoute::Pending(vec![message]));
entry.insert(messages_tx);
}
}
Ok(CellRegistration {
peer: Arc::clone(self),
key,
messages_rx,
})
}
pub(super) fn start_cell(
self: &Arc<Self>,
registration: CellRegistration,
request_id: RequestId,
started: StartedCell,
) {
debug_assert!(Arc::ptr_eq(self, &registration.peer));
debug_assert_eq!(registration.key.1, started.cell_id);
tokio::spawn(async move {
drive_cell(registration, request_id, started).await;
});
}
pub(super) fn close_cell(&self, session_id: SessionId, cell_id: CellId) {
let _ = self.route_cell_message((session_id, cell_id), CellMessage::Closed);
}
fn route_cell_message(&self, key: (SessionId, CellId), message: CellMessage) -> bool {
let routes = self
.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner);
let Some(sender) = routes.get(&key) else {
return false;
};
sender.send(message).is_ok()
}
async fn send_delegate_if_pending(
@@ -168,23 +199,28 @@ impl HostPeer {
session_id: SessionId,
request: DelegateRequest,
) {
if self.pending.lock().await.contains_key(&id) {
self.send(HostToClient::DelegateRequest {
id,
session_id,
request,
});
// Keep the pending entry locked through enqueueing the request so a
// concurrent cancellation can only be sent after it.
let pending_calls = self.pending.lock().await;
if !pending_calls.contains_key(&id) {
return;
}
self.send(HostToClient::DelegateRequest {
id,
session_id,
request,
});
drop(pending_calls);
}
}
async fn drive_cell(
peer: Arc<HostPeer>,
key: (SessionId, CellId),
mut registration: CellRegistration,
request_id: RequestId,
started: StartedCell,
mut messages_rx: mpsc::UnboundedReceiver<CellMessage>,
) {
let peer = Arc::clone(&registration.peer);
let key = registration.key.clone();
let initial_response = started.initial_response();
tokio::pin!(initial_response);
let closed = loop {
@@ -197,7 +233,7 @@ async fn drive_cell(
});
break false;
}
message = messages_rx.recv() => match message {
message = registration.messages_rx.recv() => match message {
Some(CellMessage::Delegate { id, request }) => {
peer.send_delegate_if_pending(id, key.0.clone(), request).await;
}
@@ -215,7 +251,7 @@ async fn drive_cell(
} else {
loop {
tokio::select! {
message = messages_rx.recv() => match message {
message = registration.messages_rx.recv() => match message {
Some(CellMessage::Delegate { id, request }) => {
peer.send_delegate_if_pending(id, key.0.clone(), request).await;
}
@@ -229,12 +265,19 @@ async fn drive_cell(
session_id: key.0.clone(),
cell_id: key.1.clone().into(),
});
peer.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(&key);
}
impl Drop for CellRegistration {
fn drop(&mut self) {
self.peer
.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(&self.key);
}
}
/// Cancels a reverse delegate request if its waiting future is dropped.
struct PendingDelegateRequest {
peer: Arc<HostPeer>,
id: Option<DelegateRequestId>,
@@ -257,9 +300,14 @@ impl Drop for PendingDelegateRequest {
};
let peer = Arc::clone(&self.peer);
tokio::spawn(async move {
if peer.pending.lock().await.remove(&id).is_some() {
let mut pending_calls = peer.pending.lock().await;
if pending_calls.remove(&id).is_some() {
peer.send(HostToClient::CancelDelegateRequest { id });
}
});
}
}
#[cfg(test)]
#[path = "peer_tests.rs"]
mod tests;
codex-rs/code-mode-host/src/peer_tests.rs
+68
View File
@@ -0,0 +1,68 @@
use std::time::Duration;
use codex_code_mode_protocol::RuntimeResponse;
use pretty_assertions::assert_eq;
use super::*;
#[tokio::test]
async fn registration_queues_cell_closure_before_driver_starts() {
let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded_channel();
let peer = Arc::new(HostPeer::new(outgoing_tx));
let session_id = SessionId::new("session-1").expect("session ID");
let cell_id = CellId::new("cell-1".to_string());
let request_id = RequestId::new(7);
let registration = peer
.register_cell(session_id.clone(), cell_id.clone())
.expect("register cell");
peer.close_cell(session_id.clone(), cell_id.clone());
assert!(outgoing_rx.try_recv().is_err());
let (initial_tx, initial_rx) = oneshot::channel();
peer.start_cell(
registration,
request_id,
StartedCell::from_result_receiver(cell_id.clone(), initial_rx),
);
let response = RuntimeResponse::Result {
cell_id: cell_id.clone(),
content_items: Vec::new(),
error_text: None,
};
initial_tx
.send(Ok(response.clone()))
.expect("send initial response");
assert_eq!(
outgoing_rx.recv().await,
Some(HostToClient::InitialResponse {
id: request_id,
result: WireResult::Ok {
value: response.into(),
},
})
);
assert_eq!(
outgoing_rx.recv().await,
Some(HostToClient::CellClosed {
session_id,
cell_id: cell_id.into(),
})
);
tokio::time::timeout(Duration::from_secs(1), async {
loop {
if peer
.cell_routes
.lock()
.unwrap_or_else(PoisonError::into_inner)
.is_empty()
{
break;
}
tokio::task::yield_now().await;
}
})
.await
.expect("cell route cleanup");
}
codex-rs/code-mode-protocol/src/host/message.rs
+8
View File
@@ -172,6 +172,7 @@ pub enum HostToClient {
},
}
/// An operation requested by the client after connection negotiation.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields, tag = "method", rename_all_fields = "camelCase")]
pub enum HostRequest {
@@ -197,6 +198,10 @@ pub enum HostRequest {
ShutdownSession { session_id: SessionId },
}
/// The host's immediate response to a client operation.
///
/// An accepted execution returns its initial runtime output separately as a
/// [`HostToClient::InitialResponse`].
#[derive(Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
pub enum HostResponse {
@@ -210,6 +215,7 @@ pub enum HostResponse {
SessionClosed { session_id: SessionId },
}
/// A reverse call from a hosted runtime to its client-side delegate.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
pub enum DelegateRequest {
@@ -223,6 +229,7 @@ pub enum DelegateRequest {
},
}
/// The client-side delegate's result for a reverse call from the host.
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields, tag = "type", rename_all_fields = "camelCase")]
pub enum DelegateResponse {
@@ -232,6 +239,7 @@ pub enum DelegateResponse {
NotificationDelivered,
}
/// A serializable success or error payload used by protocol response frames.
#[derive(Debug, Deserialize, PartialEq, Serialize)]
#[serde(deny_unknown_fields, tag = "status", rename_all_fields = "camelCase")]
pub enum WireResult<T> {
codex-rs/code-mode/src/remote_session.rs
+9 -1
View File
@@ -76,6 +76,10 @@ impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider {
}
}
/// Lazily maintains the shared sidecar connection used by a session provider.
///
/// It serializes process creation and allocates session IDs outside any one
/// connection so those IDs remain unique when the sidecar is replaced.
struct OwnedProcessHost {
host_program: PathBuf,
connection: StdMutex<Option<Arc<Connection>>>,
@@ -132,13 +136,17 @@ impl OwnedProcessHost {
}
}
/// Tracks a logical session's binding to the current sidecar connection.
enum SessionState {
New,
Open(Arc<Connection>),
Shutdown,
}
/// A logical code-mode session assigned to a process-owned host.
/// A logical code-mode session that lazily opens on a process-owned host.
///
/// It retains its client-assigned session and cell IDs if a failed sidecar must
/// be replaced, while the hosted runtime state itself starts over.
pub struct ProcessOwnedCodeModeSession {
process_host: Arc<OwnedProcessHost>,
session_id: SessionId,
codex-rs/code-mode/src/remote_session/connection.rs
+5
View File
@@ -38,6 +38,10 @@ mod stderr;
const IPC_CHANNEL_CAPACITY: usize = 128;
/// Owns one live sidecar process connection and its background IPC tasks.
///
/// Dropping the last connection handle cancels those tasks and causes the child
/// supervisor to terminate the sidecar.
pub(super) struct Connection {
state: Arc<ConnectionState>,
cancellation: CancellationToken,
@@ -327,6 +331,7 @@ impl Drop for Connection {
}
}
/// Shares response routing, delegates, and terminal failure state across IPC tasks.
struct ConnectionState {
outgoing_tx: mpsc::Sender<ClientToHost>,
pending: Mutex<HashMap<RequestId, oneshot::Sender<Result<HostResponse, String>>>>,
codex-rs/code-mode/src/remote_session/connection/reader.rs
+5 -1
View File
@@ -78,7 +78,7 @@ async fn handle_host_message(state: Arc<ConnectionState>, message: HostToClient)
.await;
return;
};
let cancellation = CancellationToken::new();
let cancellation = state.cancellation.child_token();
state
.delegate_cancellations
.lock()
@@ -129,3 +129,7 @@ async fn handle_host_message(state: Arc<ConnectionState>, message: HostToClient)
}
}
}
#[cfg(test)]
#[path = "reader_tests.rs"]
mod tests;
codex-rs/code-mode/src/remote_session/connection/reader_tests.rs
+97
View File
@@ -0,0 +1,97 @@
use std::sync::Mutex as StdMutex;
use std::time::Duration;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeNestedToolCall;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::NotificationFuture;
use codex_code_mode_protocol::ToolInvocationFuture;
use codex_code_mode_protocol::host::DelegateRequestId;
use codex_code_mode_protocol::host::SessionId;
use pretty_assertions::assert_eq;
use tokio::sync::oneshot;
use super::*;
struct CancellationRecordingDelegate {
cancellation_tx: StdMutex<Option<oneshot::Sender<()>>>,
}
impl CodeModeSessionDelegate for CancellationRecordingDelegate {
fn invoke_tool<'a>(
&'a self,
_invocation: CodeModeNestedToolCall,
_cancellation_token: CancellationToken,
) -> ToolInvocationFuture<'a> {
Box::pin(async { Err("unexpected tool invocation".to_string()) })
}
fn notify<'a>(
&'a self,
_call_id: String,
_cell_id: CellId,
_text: String,
cancellation_token: CancellationToken,
) -> NotificationFuture<'a> {
let cancellation_tx = self
.cancellation_tx
.lock()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.take();
Box::pin(async move {
cancellation_token.cancelled().await;
if let Some(cancellation_tx) = cancellation_tx {
let _ = cancellation_tx.send(());
}
Ok(())
})
}
fn cell_closed(&self, _cell_id: &CellId) {}
}
#[tokio::test]
async fn delegate_started_after_connection_cancellation_is_cancelled() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(/*max_capacity*/ 1);
let connection_cancellation = CancellationToken::new();
connection_cancellation.cancel();
let state = Arc::new(ConnectionState::new(outgoing_tx, connection_cancellation));
let session_id = SessionId::new("session-1").expect("session ID");
let delegate_request_id = DelegateRequestId::new(7);
let (cancellation_tx, cancellation_rx) = oneshot::channel();
state.delegates.lock().await.insert(
session_id.clone(),
Arc::new(CancellationRecordingDelegate {
cancellation_tx: StdMutex::new(Some(cancellation_tx)),
}),
);
handle_host_message(
Arc::clone(&state),
HostToClient::DelegateRequest {
id: delegate_request_id,
session_id,
request: DelegateRequest::Notify {
call_id: "call-1".to_string(),
cell_id: CellId::new("cell-1".to_string()).into(),
text: "hello".to_string(),
},
},
)
.await;
tokio::time::timeout(Duration::from_secs(1), cancellation_rx)
.await
.expect("delegate cancellation timeout")
.expect("delegate cancellation signal");
assert_eq!(
outgoing_rx.recv().await,
Some(ClientToHost::DelegateResponse {
id: delegate_request_id,
result: WireResult::Ok {
value: DelegateResponse::NotificationDelivered,
},
})
);
assert!(state.delegate_cancellations.lock().await.is_empty());
}