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

Scope code mode callback delegates to individual executions (#44865)

Browse Source 
## Why

A delegate bound at session creation cannot provide different callbacks for cells sharing that session. Each execution needs to retain its own delegate across yields and release it when the cell is cleaned up.

## What changed

- Pass `CodeModeSessionDelegate` to `execute` instead of session creation.
- Route tool calls, notifications, and cell closure callbacks through the execution's delegate in the in-process runtime and the gRPC and stdio transports.
- Retain delegates with pending executions and live cells, releasing them through closure and cancellation cleanup.

## Testing

Add coverage for distinct delegates across yielded cells, gRPC callbacks before cell admission, and delegate release after completion or abandoned execution cleanup. Update transport tests to verify callbacks reach the owning cell's delegate.

GitOrigin-RevId: 7469f52104b993e790a40c65fb800ad02b7fd606
This commit is contained in:
Abhinav
2026-09-11 17:00:07 +00:00
committed by copyberry
parent bc5957eac9
commit 3305c4f31d
30 changed files with 1027 additions and 637 deletions
Download Patch File Download Diff File Expand all files Collapse all files
codex-rs/code-mode/src/grpc_session/callbacks.rs
+12 -14
View File
@@ -133,12 +133,12 @@ impl SessionInner {
}
let invocation_id = call.invocation_id.clone();
let cancellation = match admission {
CallbackAdmission::Active(cancellation) => Ok(cancellation),
CallbackAdmission::Active(cancellation, delegate) => Ok((cancellation, delegate)),
CallbackAdmission::Cancelled => return Ok(()),
CallbackAdmission::Closed => Err(format!("code-mode cell {} is closed", call.cell_id)),
CallbackAdmission::Rejected(error) => Err(error),
};
let cancellation = match cancellation {
let (cancellation, delegate) = match cancellation {
Ok(cancellation) => cancellation,
Err(error) => {
let inner = Arc::clone(self);
@@ -157,19 +157,18 @@ impl SessionInner {
let result = match invocation {
Ok(invocation) => {
let callback = AssertUnwindSafe(async {
inner
.delegate
delegate
.invoke_tool(invocation, cancellation.child_token())
.await
})
.catch_unwind();
tokio::select! {
biased;
_ = cancellation.cancelled() => return,
result = callback => match result {
Ok(result) => result,
Err(_) => Err("code-mode tool delegate panicked".to_string()),
},
biased;
_ = cancellation.cancelled() => return,
result = callback => match result {
Ok(result) => result,
Err(_) => Err("code-mode tool delegate panicked".to_string()),
},
}
}
Err(error) => Err(error),
@@ -223,8 +222,8 @@ impl SessionInner {
.lock()
.unwrap_or_else(PoisonError::into_inner)
.admit_notification(&notification)?;
let cancellation = match admission {
CallbackAdmission::Active(cancellation) => cancellation,
let (cancellation, delegate) = match admission {
CallbackAdmission::Active(cancellation, delegate) => (cancellation, delegate),
CallbackAdmission::Cancelled | CallbackAdmission::Closed => return Ok(()),
CallbackAdmission::Rejected(error) => {
warn!("code-mode notification was dropped: {error}");
@@ -237,8 +236,7 @@ impl SessionInner {
// them without waiting for arbitrary delegate work to complete.
tokio::spawn(async move {
let callback = AssertUnwindSafe(async {
inner
.delegate
delegate
.notify(
notification.call_id,
CellId::new(notification.cell_id),
codex-rs/code-mode/src/grpc_session/mod.rs
+9 -18
View File
@@ -30,6 +30,7 @@ use tonic::transport::Channel;
use tracing::Instrument;
use self::operations::WaitSlot;
use self::state::ClosedCell;
use self::state::SessionState;
use self::transport::GrpcTransport;
use self::transport::SharedTransport;
@@ -96,7 +97,6 @@ impl GrpcCodeModeSessionProvider {
#[tracing::instrument(name = "code_mode.grpc.open_binding", level = "info", skip_all)]
async fn open_binding(
&self,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
) -> Result<Arc<GrpcCodeModeSession>, String> {
let mut client = deadline::startup("transport connection", self.transport.client()).await?;
@@ -136,7 +136,6 @@ impl GrpcCodeModeSessionProvider {
let inner = Arc::new(SessionInner {
id: opened.session_id,
client,
delegate,
runtime: tokio::runtime::Handle::current(),
state: Mutex::new(SessionState::default()),
wait_slots: Mutex::new(HashMap::new()),
@@ -180,24 +179,16 @@ impl GrpcCodeModeSessionProvider {
}
impl CodeModeSessionProvider for GrpcCodeModeSessionProvider {
fn create_session<'a>(
&'a self,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a> {
self.create_session_with_limits(delegate, CodeModeSessionCellExecutionLimits::default())
fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default())
}
fn create_session_with_limits<'a>(
&'a self,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
) -> CodeModeSessionProviderFuture<'a> {
Box::pin(async move {
let session = Arc::new(reconnect::ReconnectableSession::new(
self.clone(),
delegate,
limits,
));
let session = Arc::new(reconnect::ReconnectableSession::new(self.clone(), limits));
session.initialize().await?;
Ok(session as _)
})
@@ -229,8 +220,9 @@ impl CodeModeSession for GrpcCodeModeSession {
fn execute<'a>(
&'a self,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionResultFuture<'a, StartedCell> {
Box::pin(self.inner.execute(request))
Box::pin(self.inner.execute(request, delegate))
}
fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
@@ -255,7 +247,6 @@ impl Drop for GrpcCodeModeSession {
pub(super) struct SessionInner {
pub(super) id: String,
pub(super) client: GrpcClient,
pub(super) delegate: Arc<dyn CodeModeSessionDelegate>,
runtime: tokio::runtime::Handle,
state: Mutex<SessionState>,
wait_slots: Mutex<HashMap<CellId, Weak<WaitSlot>>>,
@@ -282,14 +273,14 @@ impl SessionInner {
Ok(())
}
pub(super) fn report_closed_cell(&self, cell_id: Option<CellId>) {
if let Some(cell_id) = cell_id {
pub(super) fn report_closed_cell(&self, cell: Option<ClosedCell>) {
if let Some((cell_id, delegate)) = cell {
self.wait_slots
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove(&cell_id);
let _ = std::panic::catch_unwind(AssertUnwindSafe(|| {
self.delegate.cell_closed(&cell_id);
delegate.cell_closed(&cell_id);
}));
}
}
codex-rs/code-mode/src/grpc_session/operations.rs
+6 -3
View File
@@ -5,6 +5,7 @@ use std::sync::atomic::Ordering;
use std::time::Duration;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::DEFAULT_EXEC_YIELD_TIME_MS;
use codex_code_mode_protocol::ExecuteRequest;
use codex_code_mode_protocol::RuntimeResponse;
@@ -45,10 +46,11 @@ impl Drop for ExecutionOwnership {
.lock()
.unwrap_or_else(PoisonError::into_inner)
.remove_execution(&self.execution_id);
let Some(cell_id) = cell else {
let Some((cell_id, delegate)) = cell else {
return;
};
self.session.report_closed_cell(Some(cell_id.clone()));
self.session
.report_closed_cell(Some((cell_id.clone(), delegate)));
if self.session.stopped.is_cancelled() {
return;
}
@@ -67,6 +69,7 @@ impl SessionInner {
pub(super) async fn execute(
self: &Arc<Self>,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> Result<StartedCell, String> {
self.require_open()?;
let execution_id = Uuid::new_v4().to_string();
@@ -81,7 +84,7 @@ impl SessionInner {
self.state
.lock()
.unwrap_or_else(PoisonError::into_inner)
.begin_execution(&request)?;
.begin_execution(&request, delegate)?;
let ownership = ExecutionOwnership {
session: Arc::clone(self),
execution_id,
codex-rs/code-mode/src/grpc_session/reconnect.rs
+7 -9
View File
@@ -31,7 +31,6 @@ pub(super) struct ReconnectableSession {
struct ReconnectInner {
provider: GrpcCodeModeSessionProvider,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
binding: Mutex<Option<SessionBinding>>,
opening_permit: Semaphore,
@@ -49,13 +48,11 @@ struct SessionBinding {
impl ReconnectableSession {
pub(super) fn new(
provider: GrpcCodeModeSessionProvider,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
) -> Self {
Self {
inner: Arc::new(ReconnectInner {
provider,
delegate,
limits,
binding: Mutex::new(None),
opening_permit: Semaphore::new(/*permits*/ 1),
@@ -75,10 +72,15 @@ impl CodeModeSession for ReconnectableSession {
fn execute<'a>(
&'a self,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionResultFuture<'a, StartedCell> {
Box::pin(async move {
let binding = self.inner.get_or_open_binding().await?;
let started = binding.session.execute(request).await?;
let delegate = Arc::new(GenerationDelegate {
delegate,
generation: binding.generation,
});
let started = binding.session.execute(request, delegate).await?;
Ok(generation::public_started_cell(binding.generation, started))
})
}
@@ -151,16 +153,12 @@ impl ReconnectInner {
}
let generation = self.next_generation.fetch_add(1, Ordering::Relaxed);
let delegate = Arc::new(GenerationDelegate {
delegate: Arc::clone(&self.delegate),
generation,
});
let session = tokio::select! {
biased;
_ = self.shutdown_requested.cancelled() => {
return Err(SHUTDOWN_ERROR.to_string());
}
session = self.provider.open_binding(delegate, self.limits.clone()) => session?,
session = self.provider.open_binding(self.limits.clone()) => session?,
};
let binding = SessionBinding {
session,
codex-rs/code-mode/src/grpc_session/state.rs
+36 -13
View File
@@ -1,6 +1,9 @@
use std::collections::HashMap;
use std::collections::HashSet;
use std::collections::VecDeque;
use std::sync::Arc;
use codex_code_mode_protocol::CodeModeSessionDelegate;
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::grpc;
@@ -16,15 +19,17 @@ struct ActiveCallback {
cancellation: CancellationToken,
}
pub(super) type ClosedCell = (CellId, Arc<dyn CodeModeSessionDelegate>);
pub(super) enum CallbackAdmission {
Active(CancellationToken),
Active(CancellationToken, Arc<dyn CodeModeSessionDelegate>),
Cancelled,
Closed,
Rejected(String),
}
#[derive(Default)]
struct ExecutionRecord {
delegate: Arc<dyn CodeModeSessionDelegate>,
cell_id: Option<CellId>,
tool_call_id: String,
enabled_tools: HashMap<ToolName, i32>,
@@ -101,7 +106,11 @@ impl SessionState {
Ok(())
}
pub(super) fn begin_execution(&mut self, request: &grpc::ExecuteRequest) -> Result<(), String> {
pub(super) fn begin_execution(
&mut self,
request: &grpc::ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> Result<(), String> {
self.require_open()?;
if request.execution_id.is_empty() || self.executions.contains_key(&request.execution_id) {
return Err("code-mode execution ID was empty or reused".to_string());
@@ -127,7 +136,13 @@ impl SessionState {
ExecutionRecord {
tool_call_id: request.tool_call_id.clone(),
enabled_tools,
..ExecutionRecord::default()
delegate,
cell_id: None,
started: false,
ready: false,
closed: false,
notifications: 0,
cancellation: CancellationToken::new(),
},
);
Ok(())
@@ -155,7 +170,7 @@ impl SessionState {
pub(super) fn mark_execution_ready(
&mut self,
execution_id: &str,
) -> Result<Option<CellId>, String> {
) -> Result<Option<ClosedCell>, String> {
self.require_open()?;
let execution = self
.executions
@@ -224,7 +239,10 @@ impl SessionState {
cancellation: cancellation.clone(),
},
);
Ok(CallbackAdmission::Active(cancellation))
Ok(CallbackAdmission::Active(
cancellation,
Arc::clone(&execution.delegate),
))
}
pub(super) fn admit_notification(
@@ -255,10 +273,11 @@ impl SessionState {
self.notifications += 1;
Ok(CallbackAdmission::Active(
execution.cancellation.child_token(),
Arc::clone(&execution.delegate),
))
}
pub(super) fn finish_notification(&mut self, execution_id: &str) -> Option<CellId> {
pub(super) fn finish_notification(&mut self, execution_id: &str) -> Option<ClosedCell> {
let execution = self.executions.get_mut(execution_id)?;
execution.notifications = execution.notifications.checked_sub(1)?;
self.notifications -= 1;
@@ -296,7 +315,7 @@ impl SessionState {
pub(super) fn close_cell(
&mut self,
closed: grpc::CellClosed,
) -> Result<Option<CellId>, String> {
) -> Result<Option<ClosedCell>, String> {
self.require_open()?;
self.check_cell_ownership(&closed.execution_id, &closed.cell_id)?;
let Some(execution) = self.executions.get_mut(&closed.execution_id) else {
@@ -317,7 +336,7 @@ impl SessionState {
Ok(self.close_execution_if_ready(&closed.execution_id))
}
pub(super) fn close(&mut self, failure: Option<String>) -> Vec<CellId> {
pub(super) fn close(&mut self, failure: Option<String>) -> Vec<ClosedCell> {
if self.closed {
return Vec::new();
}
@@ -331,12 +350,14 @@ impl SessionState {
.drain()
.filter_map(|(_, execution)| {
execution.cancellation.cancel();
execution.cell_id
execution
.cell_id
.map(|cell_id| (cell_id, execution.delegate))
})
.collect()
}
fn close_execution_if_ready(&mut self, execution_id: &str) -> Option<CellId> {
fn close_execution_if_ready(&mut self, execution_id: &str) -> Option<ClosedCell> {
self.executions
.get(execution_id)
.is_some_and(|execution| {
@@ -349,12 +370,14 @@ impl SessionState {
.flatten()
}
pub(super) fn remove_execution(&mut self, execution_id: &str) -> Option<CellId> {
pub(super) fn remove_execution(&mut self, execution_id: &str) -> Option<ClosedCell> {
let execution = self.executions.remove(execution_id)?;
self.notifications -= execution.notifications;
execution.cancellation.cancel();
self.revoke_execution_callbacks(execution_id);
execution.cell_id
execution
.cell_id
.map(|cell_id| (cell_id, execution.delegate))
}
fn check_cell_ownership(&self, execution_id: &str, cell_id: &str) -> Result<(), String> {
codex-rs/code-mode/src/grpc_session/state_tests.rs
+95 -50
View File
@@ -1,7 +1,9 @@
use codex_code_mode_protocol::CellId;
use codex_code_mode_protocol::NoopCodeModeSessionDelegate;
use codex_code_mode_protocol::grpc;
use codex_code_mode_protocol::host::MAX_PENDING_DELEGATE_CALLS;
use pretty_assertions::assert_eq;
use std::sync::Arc;
use uuid::Uuid;
use super::CallbackAdmission;
@@ -58,33 +60,66 @@ fn notification(execution_id: &str, notification_id: u128) -> grpc::Notification
}
}
#[test]
fn callbacks_use_their_execution_delegate_before_cell_admission() {
let mut state = SessionState::default();
let delegates: Vec<Arc<dyn codex_code_mode_protocol::CodeModeSessionDelegate>> = vec![
Arc::new(NoopCodeModeSessionDelegate),
Arc::new(NoopCodeModeSessionDelegate),
];
for (index, delegate) in delegates.iter().enumerate() {
state
.begin_execution(&request(&index.to_string()), delegate.clone())
.unwrap();
}
for (index, delegate) in delegates.iter().enumerate() {
let mut call = tool_call(&index.to_string(), index as u128 + 1);
call.cell_id = format!("cell-{index}");
let CallbackAdmission::Active(_, actual) = state.admit_invocation(&call).unwrap() else {
panic!("callback should be admitted before the execute response");
};
assert!(Arc::ptr_eq(&actual, delegate));
let mut notice = notification(&index.to_string(), index as u128 + 1);
notice.cell_id = call.cell_id;
let CallbackAdmission::Active(_, actual) = state.admit_notification(&notice).unwrap()
else {
panic!("notification should be admitted before the execute response");
};
assert!(Arc::ptr_eq(&actual, delegate));
}
drop(state.close(/*failure*/ None));
for delegate in delegates {
assert_eq!(Arc::strong_count(&delegate), 1);
}
}
#[test]
fn cell_closure_drains_notifications_and_cancels_tool_callbacks() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
let first = state
.admit_invocation(&tool_call("execution", /*invocation_id*/ 1))
.expect("accept early invocation");
let CallbackAdmission::Active(first_cancellation) = first else {
let CallbackAdmission::Active(first_cancellation, _) = first else {
panic!("first callback was not admitted");
};
let CallbackAdmission::Active(notification_cancellation) = state
let CallbackAdmission::Active(notification_cancellation, _) = state
.admit_notification(&notification("execution", /*notification_id*/ 1))
.expect("admit notification")
else {
panic!("notification was not admitted");
};
assert_eq!(
assert!(
state
.close_cell(grpc::CellClosed {
execution_id: "execution".to_string(),
cell_id: "cell".to_string(),
final_tool_call_sequence: 3,
})
.expect("record cell closure"),
None
.expect("record cell closure")
.is_none()
);
assert!(first_cancellation.is_cancelled());
assert!(!notification_cancellation.is_cancelled());
@@ -97,14 +132,16 @@ fn cell_closure_drains_notifications_and_cancels_tool_callbacks() {
.expect("reject invocation for a closed cell"),
CallbackAdmission::Closed
));
assert_eq!(
assert!(
state
.mark_execution_ready("execution")
.expect("claim started cell"),
None
.expect("claim started cell")
.is_none()
);
assert_eq!(
state.finish_notification("execution"),
state
.finish_notification("execution")
.map(|(cell_id, _)| cell_id),
Some(CellId::new("cell".to_string()))
);
assert!(notification_cancellation.is_cancelled());
@@ -114,18 +151,18 @@ fn cell_closure_drains_notifications_and_cancels_tool_callbacks() {
fn cell_closure_waits_until_the_started_cell_is_claimed() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
assert_eq!(
assert!(
state
.close_cell(grpc::CellClosed {
execution_id: "execution".to_string(),
cell_id: "cell".to_string(),
final_tool_call_sequence: 0,
})
.expect("record early cell closure"),
None
.expect("record early cell closure")
.is_none()
);
state
.admit_execution("execution", "cell")
@@ -133,6 +170,7 @@ fn cell_closure_waits_until_the_started_cell_is_claimed() {
assert_eq!(
state
.mark_execution_ready("execution")
.map(|cell| cell.map(|(cell_id, _)| cell_id))
.expect("claim started cell"),
Some(CellId::new("cell".to_string()))
);
@@ -142,7 +180,7 @@ fn cell_closure_waits_until_the_started_cell_is_claimed() {
fn oversized_cell_ids_are_rejected_before_admission() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
assert_eq!(
@@ -152,17 +190,17 @@ fn oversized_cell_ids_are_rejected_before_admission() {
grpc::MAX_IDENTIFIER_BYTES
))
);
assert_eq!(state.remove_execution("execution"), None);
assert!(state.remove_execution("execution").is_none());
}
#[test]
fn callbacks_cannot_claim_another_executions_cell() {
let mut state = SessionState::default();
state
.begin_execution(&request("first"))
.begin_execution(&request("first"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register first execution");
state
.begin_execution(&request("second"))
.begin_execution(&request("second"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register second execution");
state
.admit_invocation(&tool_call("first", /*invocation_id*/ 1))
@@ -198,19 +236,19 @@ fn callbacks_cannot_claim_another_executions_cell() {
fn abandonment_before_start_ignores_later_cell_closure() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
assert_eq!(state.remove_execution("execution"), None);
assert_eq!(
assert!(state.remove_execution("execution").is_none());
assert!(
state
.close_cell(grpc::CellClosed {
execution_id: "execution".to_string(),
cell_id: "cell".to_string(),
final_tool_call_sequence: 0,
})
.expect("ignore closure for abandoned execution"),
None
.expect("ignore closure for abandoned execution")
.is_none()
);
assert!(state.close(/*failure*/ None).is_empty());
}
@@ -219,16 +257,18 @@ fn abandonment_before_start_ignores_later_cell_closure() {
fn abandonment_revokes_callbacks_and_ignores_late_events() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
let CallbackAdmission::Active(invocation) = state
let CallbackAdmission::Active(invocation, _) = state
.admit_invocation(&tool_call("execution", /*invocation_id*/ 1))
.expect("admit callback before execution starts")
else {
panic!("invocation was not admitted");
};
assert_eq!(
state.remove_execution("execution"),
state
.remove_execution("execution")
.map(|(cell_id, _)| cell_id),
Some(CellId::new("cell".to_string()))
);
assert!(invocation.is_cancelled());
@@ -238,15 +278,15 @@ fn abandonment_revokes_callbacks_and_ignores_late_events() {
.expect("reject delayed tool invocation"),
CallbackAdmission::Closed
));
assert_eq!(
assert!(
state
.close_cell(grpc::CellClosed {
execution_id: "execution".to_string(),
cell_id: "cell".to_string(),
final_tool_call_sequence: 1,
})
.expect("ignore delayed cell closure"),
None
.expect("ignore delayed cell closure")
.is_none()
);
assert!(state.close(/*failure*/ None).is_empty());
}
@@ -255,7 +295,7 @@ fn abandonment_revokes_callbacks_and_ignores_late_events() {
fn invocation_cancellation_revokes_delegate_and_late_completion() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
state
.admit_execution("execution", "cell")
@@ -264,7 +304,7 @@ fn invocation_cancellation_revokes_delegate_and_late_completion() {
.mark_execution_ready("execution")
.expect("claim execution");
let invocation = tool_call("execution", /*invocation_id*/ 1);
let CallbackAdmission::Active(cancellation) = state
let CallbackAdmission::Active(cancellation, _) = state
.admit_invocation(&invocation)
.expect("accept invocation")
else {
@@ -284,6 +324,7 @@ fn invocation_cancellation_revokes_delegate_and_late_completion() {
cell_id: "cell".to_string(),
final_tool_call_sequence: 1,
})
.map(|cell| cell.map(|(cell_id, _)| cell_id))
.expect("close cell"),
Some(CellId::new("cell".to_string()))
);
@@ -293,7 +334,7 @@ fn invocation_cancellation_revokes_delegate_and_late_completion() {
fn duplicate_invocation_ids_are_rejected() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
state
.admit_execution("execution", "cell")
@@ -311,7 +352,7 @@ fn duplicate_invocation_ids_are_rejected() {
fn tool_callbacks_must_match_the_executions_enabled_tools() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
let mut disabled = tool_call("execution", /*invocation_id*/ 1);
@@ -349,7 +390,7 @@ fn tool_callbacks_must_match_the_executions_enabled_tools() {
});
assert!(matches!(
state.admit_invocation(&explicit_default_namespace),
Ok(CallbackAdmission::Active(_))
Ok(CallbackAdmission::Active(_, _))
));
assert_eq!(state.require_open(), Ok(()));
}
@@ -361,20 +402,20 @@ fn execution_call_ids_must_be_bounded() {
oversized.tool_call_id = "x".repeat(grpc::MAX_IDENTIFIER_BYTES + 1);
assert_eq!(
state.begin_execution(&oversized),
state.begin_execution(&oversized, Arc::new(NoopCodeModeSessionDelegate)),
Err(format!(
"gRPC code-mode host returned tool call ID exceeding {} bytes",
grpc::MAX_IDENTIFIER_BYTES
))
);
assert_eq!(state.remove_execution("execution"), None);
assert!(state.remove_execution("execution").is_none());
}
#[test]
fn notification_call_ids_must_match_their_execution() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
let mut oversized = notification("execution", /*notification_id*/ 1);
@@ -395,7 +436,7 @@ fn notification_call_ids_must_match_their_execution() {
);
assert!(matches!(
state.admit_notification(&notification("execution", /*notification_id*/ 3)),
Ok(CallbackAdmission::Active(_))
Ok(CallbackAdmission::Active(_, _))
));
}
@@ -403,7 +444,7 @@ fn notification_call_ids_must_match_their_execution() {
fn malformed_callback_ids_are_rejected_before_retention() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
let mut invalid_invocation = tool_call("execution", /*invocation_id*/ 1);
invalid_invocation.invocation_id = "not-a-uuid".to_string();
@@ -442,13 +483,13 @@ fn malformed_callback_ids_are_rejected_before_retention() {
fn notifications_and_tools_share_the_pending_delegate_limit() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
for index in 0..MAX_PENDING_DELEGATE_CALLS {
assert!(matches!(
state.admit_notification(&notification("execution", index as u128 + 1)),
Ok(CallbackAdmission::Active(_))
Ok(CallbackAdmission::Active(_, _))
));
}
@@ -463,10 +504,10 @@ fn notifications_and_tools_share_the_pending_delegate_limit() {
if error == "code-mode host exceeded its pending delegate callback limit"
));
assert_eq!(state.require_open(), Ok(()));
assert_eq!(state.finish_notification("execution"), None);
assert!(state.finish_notification("execution").is_none());
assert!(matches!(
state.admit_invocation(&tool_call("execution", /*invocation_id*/ 2)),
Ok(CallbackAdmission::Active(_))
Ok(CallbackAdmission::Active(_, _))
));
}
@@ -474,9 +515,9 @@ fn notifications_and_tools_share_the_pending_delegate_limit() {
fn terminated_cells_cancel_pending_notifications() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
let CallbackAdmission::Active(cancellation) = state
let CallbackAdmission::Active(cancellation, _) = state
.admit_notification(&notification("execution", /*notification_id*/ 1))
.expect("admit notification")
else {
@@ -486,25 +527,25 @@ fn terminated_cells_cancel_pending_notifications() {
state.cancel_notifications(&CellId::new("cell".to_string()));
assert!(cancellation.is_cancelled());
assert_eq!(state.finish_notification("execution"), None);
assert!(state.finish_notification("execution").is_none());
}
#[test]
fn disconnect_revokes_callbacks_and_returns_each_live_cell_once() {
let mut state = SessionState::default();
state
.begin_execution(&request("execution"))
.begin_execution(&request("execution"), Arc::new(NoopCodeModeSessionDelegate))
.expect("register execution");
state
.admit_execution("execution", "cell")
.expect("admit execution");
let CallbackAdmission::Active(cancellation) = state
let CallbackAdmission::Active(cancellation, _) = state
.admit_invocation(&tool_call("execution", /*invocation_id*/ 1))
.expect("accept invocation")
else {
panic!("invocation was not admitted");
};
let CallbackAdmission::Active(notification_cancellation) = state
let CallbackAdmission::Active(notification_cancellation, _) = state
.admit_notification(&notification("execution", /*notification_id*/ 1))
.expect("admit notification")
else {
@@ -512,7 +553,11 @@ fn disconnect_revokes_callbacks_and_returns_each_live_cell_once() {
};
assert_eq!(
state.close(Some("lease closed".to_string())),
state
.close(Some("lease closed".to_string()))
.into_iter()
.map(|(cell_id, _)| cell_id)
.collect::<Vec<_>>(),
vec![CellId::new("cell".to_string())]
);
assert!(cancellation.is_cancelled());
codex-rs/code-mode/src/remote_session.rs
+21 -32
View File
@@ -26,7 +26,6 @@ use self::connection::Connection;
use self::connection::ConnectionError;
use self::connection::RemoteSession;
use self::connection::SessionCleanup;
use crate::NoopCodeModeSessionDelegate;
mod connection;
@@ -73,19 +72,15 @@ impl CodeModeSessionProvider for ProcessOwnedCodeModeSessionProvider {
}
}
fn create_session<'a>(
&'a self,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a> {
self.create_session_with_limits(delegate, CodeModeSessionCellExecutionLimits::default())
fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
self.create_session_with_limits(CodeModeSessionCellExecutionLimits::default())
}
fn create_session_with_limits<'a>(
&'a self,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
) -> CodeModeSessionProviderFuture<'a> {
Box::pin(create_host_session(delegate, self.process_host(), limits))
Box::pin(create_host_session(self.process_host(), limits))
}
}
@@ -94,28 +89,23 @@ impl CodeModeSessionProvider for DisabledCodeModeSessionProvider {
Err("code-mode host is disabled".to_string())
}
fn create_session<'a>(
&'a self,
_delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionProviderFuture<'a> {
fn create_session(&self) -> CodeModeSessionProviderFuture<'_> {
Box::pin(async { Err("code-mode host is disabled".to_string()) })
}
fn create_session_with_limits<'a>(
&'a self,
delegate: Arc<dyn CodeModeSessionDelegate>,
_limits: CodeModeSessionCellExecutionLimits,
) -> CodeModeSessionProviderFuture<'a> {
self.create_session(delegate)
self.create_session()
}
}
async fn create_host_session(
delegate: Arc<dyn CodeModeSessionDelegate>,
host: Arc<OwnedCodeModeHost>,
limits: CodeModeSessionCellExecutionLimits,
) -> Result<Arc<dyn CodeModeSession>, String> {
let session = ProcessOwnedCodeModeSession::with_host(delegate, host, limits);
let session = ProcessOwnedCodeModeSession::with_host(host, limits);
session.connection().await?;
Ok(Arc::new(session))
}
@@ -223,7 +213,6 @@ struct SessionBinding {
struct SessionInner {
host: Arc<OwnedCodeModeHost>,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
state: StdMutex<SessionState>,
next_generation: AtomicU64,
@@ -240,7 +229,6 @@ pub struct ProcessOwnedCodeModeSession {
impl ProcessOwnedCodeModeSession {
pub fn new() -> Self {
Self::with_host(
Arc::new(NoopCodeModeSessionDelegate),
Arc::new(OwnedCodeModeHost::new(
InstallContext::current().code_mode_host_program(),
)),
@@ -248,15 +236,10 @@ impl ProcessOwnedCodeModeSession {
)
}
fn with_host(
delegate: Arc<dyn CodeModeSessionDelegate>,
host: Arc<OwnedCodeModeHost>,
limits: CodeModeSessionCellExecutionLimits,
) -> Self {
fn with_host(host: Arc<OwnedCodeModeHost>, limits: CodeModeSessionCellExecutionLimits) -> Self {
Self {
inner: Arc::new(SessionInner {
host,
delegate,
limits,
state: StdMutex::new(SessionState::New),
next_generation: AtomicU64::new(1),
@@ -271,9 +254,16 @@ impl ProcessOwnedCodeModeSession {
self.inner.connection().await
}
pub async fn execute(&self, request: ExecuteRequest) -> Result<StartedCell, String> {
pub async fn execute(
&self,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> Result<StartedCell, String> {
let binding = self.connection().await?;
binding.connection.execute(binding.remote, request).await
binding
.connection
.execute(binding.remote, request, delegate)
.await
}
pub async fn wait(&self, request: WaitRequest) -> Result<WaitOutcome, String> {
@@ -348,11 +338,7 @@ impl SessionInner {
let result = match self.host.connection().await {
Ok(connection) => {
let cleanup = connection
.open_session(
remote.clone(),
Arc::clone(&self.delegate),
self.limits.clone(),
)
.open_session(remote.clone(), self.limits.clone())
.await;
cleanup.map(|cleanup| SessionBinding {
connection,
@@ -532,8 +518,11 @@ impl CodeModeSession for ProcessOwnedCodeModeSession {
fn execute<'a>(
&'a self,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> CodeModeSessionResultFuture<'a, StartedCell> {
Box::pin(ProcessOwnedCodeModeSession::execute(self, request))
Box::pin(ProcessOwnedCodeModeSession::execute(
self, request, delegate,
))
}
fn wait<'a>(&'a self, request: WaitRequest) -> CodeModeSessionResultFuture<'a, WaitOutcome> {
codex-rs/code-mode/src/remote_session/connection.rs
+2 -2
View File
@@ -323,7 +323,6 @@ impl Connection {
pub(super) async fn open_session(
&self,
session: RemoteSession,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
) -> Result<SessionCleanup, String> {
if limits != CodeModeSessionCellExecutionLimits::default()
@@ -341,7 +340,6 @@ impl Connection {
let (response_tx, response_rx) = oneshot::channel();
self.send(DriverCommand::OpenSession {
session,
delegate,
limits,
cleanup: cleanup.clone(),
caller_cancellation: cancellation.token(),
@@ -358,12 +356,14 @@ impl Connection {
&self,
session: RemoteSession,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> Result<StartedCell, String> {
let cancellation = CallerCancellation::new();
let (response_tx, response_rx) = oneshot::channel();
self.send(DriverCommand::Execute {
session,
request,
delegate,
caller_cancellation: cancellation.token(),
response_tx,
})
codex-rs/code-mode/src/remote_session/connection/driver/commands.rs
+5 -12
View File
@@ -29,25 +29,18 @@ impl ConnectionDriver {
match command {
DriverCommand::OpenSession {
session,
delegate,
limits,
cleanup,
caller_cancellation,
response_tx,
} => self.open_session(
session,
delegate,
limits,
cleanup,
caller_cancellation,
response_tx,
),
} => self.open_session(session, limits, cleanup, caller_cancellation, response_tx),
DriverCommand::Execute {
session,
request,
delegate,
caller_cancellation,
response_tx,
} => self.execute(session, request, caller_cancellation, response_tx),
} => self.execute(session, request, delegate, caller_cancellation, response_tx),
DriverCommand::Wait {
session,
request,
@@ -69,7 +62,6 @@ impl ConnectionDriver {
fn open_session(
&mut self,
session: RemoteSession,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
cleanup: super::cleanup::SessionCleanup,
caller_cancellation: CancellationToken,
@@ -120,7 +112,6 @@ impl ConnectionDriver {
request_id,
PendingRequest::OpenSession {
session,
delegate,
cleanup,
cancellation,
response_tx,
@@ -134,6 +125,7 @@ impl ConnectionDriver {
&mut self,
session: RemoteSession,
request: ExecuteRequest,
delegate: Arc<dyn CodeModeSessionDelegate>,
caller_cancellation: CancellationToken,
response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
) -> bool {
@@ -179,6 +171,7 @@ impl ConnectionDriver {
request_id,
PendingRequest::Execute {
session,
delegate,
response_tx,
initial_response_tx,
initial_response_rx,
codex-rs/code-mode/src/remote_session/connection/driver/responses.rs
+3 -4
View File
@@ -153,15 +153,13 @@ impl ConnectionDriver {
match pending {
PendingRequest::OpenSession {
session,
delegate,
cleanup,
cancellation,
response_tx,
} => match result {
Ok(HostResponse::SessionReady { session_id }) if session_id == session.id => {
let abandoned = cancellation.is_cancelled() || response_tx.is_closed();
self.sessions
.insert_ready(session.clone(), delegate, cleanup);
self.sessions.insert_ready(session.clone(), cleanup);
if abandoned || response_tx.send(Ok(())).is_err() {
return self.shutdown_abandoned_session(session);
}
@@ -179,6 +177,7 @@ impl ConnectionDriver {
},
PendingRequest::Execute {
session,
delegate,
response_tx,
initial_response_tx,
initial_response_rx,
@@ -188,7 +187,7 @@ impl ConnectionDriver {
// The host owns a checked, never-reused ID sequence. Retain only live
// IDs so client memory scales with concurrency, not session lifetime.
let remote_cell_id = cell_id.clone();
let public_id = match self.sessions.admit_cell(&session, cell_id) {
let public_id = match self.sessions.admit_cell(&session, cell_id, delegate) {
Ok(public_id) => public_id,
Err(CellAdmissionError::MissingSession) => {
let _ = response_tx
codex-rs/code-mode/src/remote_session/connection/driver/session_registry.rs
+20 -40
View File
@@ -34,10 +34,9 @@ pub(super) enum CellAdmissionError {
struct SessionRecord {
remote: RemoteSession,
delegate: Arc<dyn CodeModeSessionDelegate>,
cleanup: SessionCleanup,
phase: SessionPhase,
cells: HashMap<WireCellId, CellId>,
cells: HashMap<WireCellId, CellOwner>,
}
#[derive(Clone, Copy, Eq, PartialEq)]
@@ -67,17 +66,11 @@ impl SessionRegistry {
.is_some_and(|session| session.cells.contains_key(cell_id))
}
pub(super) fn insert_ready(
&mut self,
session: RemoteSession,
delegate: Arc<dyn CodeModeSessionDelegate>,
cleanup: SessionCleanup,
) {
pub(super) fn insert_ready(&mut self, session: RemoteSession, cleanup: SessionCleanup) {
self.records.insert(
session.id.clone(),
SessionRecord {
remote: session,
delegate,
cleanup,
phase: SessionPhase::Ready,
cells: HashMap::new(),
@@ -133,6 +126,7 @@ impl SessionRegistry {
&mut self,
session: &RemoteSession,
cell_id: WireCellId,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> Result<CellId, CellAdmissionError> {
let Some(record) = self.records.get_mut(&session.id) else {
return Err(CellAdmissionError::MissingSession);
@@ -141,7 +135,14 @@ impl SessionRegistry {
return Err(CellAdmissionError::DuplicateCell);
}
let public_id = public_cell_id(session.generation, &cell_id);
record.cells.insert(cell_id, public_id.clone());
record.cells.insert(
cell_id,
CellOwner {
session_id: session.id.clone(),
cell_id: public_id.clone(),
delegate,
},
);
Ok(public_id)
}
@@ -154,7 +155,7 @@ impl SessionRegistry {
.records
.get(session_id)
.ok_or_else(|| format!("code-mode host delegated for unknown session {session_id}"))?;
let public_id = session.cells.get(cell_id).cloned().ok_or_else(|| {
let owner = session.cells.get(cell_id).ok_or_else(|| {
format!(
"code-mode host delegated for unknown cell {} in session {session_id}",
cell_id.as_str()
@@ -162,8 +163,8 @@ impl SessionRegistry {
})?;
Ok(DelegateTarget {
session_id: session_id.clone(),
cell_id: public_id,
delegate: Arc::clone(&session.delegate),
cell_id: owner.cell_id.clone(),
delegate: Arc::clone(&owner.delegate),
})
}
@@ -178,46 +179,25 @@ impl SessionRegistry {
cell_id.as_str()
)
})?;
let public_id = session
session
.cells
.remove(cell_id)
.ok_or_else(|| format!("code-mode host closed unknown cell in session {session_id}"))?;
Ok(CellOwner {
session_id: session_id.clone(),
cell_id: public_id,
delegate: Arc::clone(&session.delegate),
})
.ok_or_else(|| format!("code-mode host closed unknown cell in session {session_id}"))
}
pub(super) fn remove_session(&mut self, session_id: &SessionId) -> Vec<CellOwner> {
let Some(session) = self.records.remove(session_id) else {
return Vec::new();
};
session
.cells
.into_values()
.map(|cell_id| CellOwner {
session_id: session_id.clone(),
cell_id,
delegate: Arc::clone(&session.delegate),
})
.collect()
session.cells.into_values().collect()
}
pub(super) fn drain(&mut self) -> Vec<FailedSession> {
let sessions = std::mem::take(&mut self.records);
sessions
.into_iter()
.map(|(session_id, session)| {
let cells = session
.cells
.into_values()
.map(|cell_id| CellOwner {
session_id: session_id.clone(),
cell_id,
delegate: Arc::clone(&session.delegate),
})
.collect();
.into_values()
.map(|session| {
let cells = session.cells.into_values().collect();
FailedSession {
cleanup: session.cleanup,
cells,
codex-rs/code-mode/src/remote_session/connection/driver/types.rs
+2 -2
View File
@@ -30,7 +30,6 @@ pub(in crate::remote_session) struct RemoteSession {
pub(in crate::remote_session::connection) enum DriverCommand {
OpenSession {
session: RemoteSession,
delegate: Arc<dyn CodeModeSessionDelegate>,
limits: CodeModeSessionCellExecutionLimits,
cleanup: SessionCleanup,
caller_cancellation: CancellationToken,
@@ -38,6 +37,7 @@ pub(in crate::remote_session::connection) enum DriverCommand {
},
Execute {
session: RemoteSession,
delegate: Arc<dyn CodeModeSessionDelegate>,
request: ExecuteRequest,
caller_cancellation: CancellationToken,
response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
@@ -136,13 +136,13 @@ pub(super) struct UnclaimedExecute {
pub(super) enum PendingRequest {
OpenSession {
session: RemoteSession,
delegate: Arc<dyn CodeModeSessionDelegate>,
cleanup: SessionCleanup,
cancellation: CancellableRequest,
response_tx: oneshot::Sender<Result<(), String>>,
},
Execute {
session: RemoteSession,
delegate: Arc<dyn CodeModeSessionDelegate>,
response_tx: oneshot::Sender<Result<DeliveredExecute, String>>,
initial_response_tx: oneshot::Sender<Result<RuntimeResponse, String>>,
initial_response_rx: oneshot::Receiver<Result<RuntimeResponse, String>>,
codex-rs/code-mode/src/remote_session/connection/driver_tests.rs
+163 -67
View File
@@ -87,17 +87,13 @@ impl DriverHarness {
}
}
async fn open(
&mut self,
session: RemoteSession,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> SessionCleanup {
async fn open(&mut self, session: RemoteSession) -> SessionCleanup {
let cleanup = SessionCleanup::new();
let (response_tx, response_rx) = oneshot::channel();
self.command_tx
.send(DriverCommand::OpenSession {
session: session.clone(),
delegate,
limits: Default::default(),
cleanup: cleanup.clone(),
caller_cancellation: CancellationToken::new(),
@@ -129,11 +125,13 @@ impl DriverHarness {
session: RemoteSession,
request_id: i64,
cell_id: &str,
delegate: Arc<dyn CodeModeSessionDelegate>,
) -> codex_code_mode_protocol::StartedCell {
let (response_tx, response_rx) = oneshot::channel();
self.command_tx
.send(DriverCommand::Execute {
session,
delegate,
request: ExecuteRequest {
tool_call_id: format!("call-{request_id}"),
enabled_tools: Vec::new(),
@@ -208,7 +206,7 @@ async fn open_session_includes_nondefault_cell_execution_limits() {
.command_tx
.send(DriverCommand::OpenSession {
session: session.clone(),
delegate: Arc::new(RecordingDelegate::default()),
limits,
cleanup: SessionCleanup::new(),
caller_cancellation: CancellationToken::new(),
@@ -388,18 +386,20 @@ async fn next_held_delegate_event(
async fn deferred_delegates_follow_cell_readiness_and_cancellation() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
let first_delegate = Arc::new(RecordingDelegate::default());
let second_delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone()).await;
let (first_response_tx, first_response_rx) = oneshot::channel();
let (second_response_tx, second_response_rx) = oneshot::channel();
for (tool_call_id, response_tx) in [
("first-cell", first_response_tx),
("second-cell", second_response_tx),
for (tool_call_id, response_tx, delegate) in [
("first-cell", first_response_tx, first_delegate.clone()),
("second-cell", second_response_tx, second_delegate.clone()),
] {
harness
.command_tx
.send(DriverCommand::Execute {
delegate: delegate.clone(),
session: session.clone(),
request: ExecuteRequest {
tool_call_id: tool_call_id.to_string(),
@@ -481,7 +481,13 @@ async fn deferred_delegates_follow_cell_readiness_and_cancellation() {
},
}
);
assert_eq!(delegate.notifications.load(Ordering::Relaxed), 1);
assert_eq!(
(
first_delegate.notifications.load(Ordering::Relaxed),
second_delegate.notifications.load(Ordering::Relaxed)
),
(0, 1)
);
harness
.event_tx
@@ -513,7 +519,13 @@ async fn deferred_delegates_follow_cell_readiness_and_cancellation() {
},
}
);
assert_eq!(delegate.notifications.load(Ordering::Relaxed), 2);
assert_eq!(
(
first_delegate.notifications.load(Ordering::Relaxed),
second_delegate.notifications.load(Ordering::Relaxed)
),
(1, 1)
);
assert!(matches!(
harness.outgoing_rx.try_recv(),
Err(mpsc::error::TryRecvError::Empty)
@@ -531,7 +543,7 @@ async fn dropped_open_waiter_shuts_down_committed_session() {
.command_tx
.send(DriverCommand::OpenSession {
session: session.clone(),
delegate: Arc::new(RecordingDelegate::default()),
limits: Default::default(),
cleanup,
caller_cancellation: CancellationToken::new(),
@@ -575,6 +587,7 @@ async fn dropped_open_waiter_shuts_down_committed_session() {
harness
.command_tx
.send(DriverCommand::Execute {
delegate: Arc::new(RecordingDelegate::default()),
session: session.clone(),
request: ExecuteRequest {
tool_call_id: "call-1".to_string(),
@@ -603,9 +616,14 @@ async fn delegate_cancel_is_best_effort_and_sends_no_late_response() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
let request_id = DelegateRequestId::new(/*value*/ 7);
harness
@@ -662,9 +680,14 @@ async fn delegate_limit_returns_an_error_without_disconnecting() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
for value in 1..=MAX_PENDING_DELEGATE_CALLS {
@@ -729,9 +752,14 @@ async fn terminate_closes_cell_without_waiting_for_delegate_cleanup() {
let mut harness = DriverHarness::start();
let session = remote_session();
let (delegate, mut events_rx, release) = HeldDelegate::new();
harness.open(session.clone(), delegate).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
let delegate_id = DelegateRequestId::new(/*value*/ 7);
harness.start_tool_delegate(&session, delegate_id).await;
@@ -825,9 +853,14 @@ async fn shutdown_closes_cell_without_waiting_for_delegate_cleanup() {
let mut harness = DriverHarness::start();
let session = remote_session();
let (delegate, mut events_rx, release) = HeldDelegate::new();
harness.open(session.clone(), delegate).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
let delegate_id = DelegateRequestId::new(/*value*/ 7);
harness.start_tool_delegate(&session, delegate_id).await;
@@ -907,9 +940,14 @@ async fn completed_delegate_request_id_cannot_be_reused() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
let request_id = DelegateRequestId::new(/*value*/ 7);
let request = || DelegateRequest::Notify {
@@ -950,11 +988,14 @@ async fn completed_delegate_request_id_cannot_be_reused() {
async fn delegate_task_panic_becomes_tool_error_without_killing_connection() {
let mut harness = DriverHarness::start();
let session = remote_session();
harness
.open(session.clone(), Arc::new(PanickingDelegate))
.await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
Arc::new(PanickingDelegate),
)
.await;
harness
.event_tx
@@ -994,7 +1035,15 @@ async fn delegate_for_unknown_cell_returns_error_without_invocation() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(
session.clone(),
/*request_id*/ 2,
"known",
delegate.clone(),
)
.await;
let id = DelegateRequestId::new(/*value*/ 7);
harness
@@ -1039,9 +1088,14 @@ async fn delegate_after_cell_close_returns_error_without_invocation() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
harness
.event_tx
@@ -1090,8 +1144,10 @@ async fn mismatched_initial_response_fails_connection_and_closes_cell_once() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
let started = harness.start_cell(session, /*request_id*/ 2, "1").await;
harness.open(session.clone()).await;
let started = harness
.start_cell(session, /*request_id*/ 2, "1", delegate.clone())
.await;
harness
.event_tx
.send(DriverEvent::HostMessage(HostToClient::InitialResponse {
@@ -1120,9 +1176,14 @@ async fn mismatched_wait_response_fails_connection() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
let (response_tx, response_rx) = oneshot::channel();
harness
@@ -1169,9 +1230,14 @@ async fn mismatched_terminate_response_fails_connection() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
let (response_tx, response_rx) = oneshot::channel();
harness
@@ -1213,11 +1279,14 @@ async fn mismatched_terminate_response_fails_connection() {
async fn remote_wait_accepts_durations_longer_than_five_minutes() {
let mut harness = DriverHarness::start();
let session = remote_session();
harness
.open(session.clone(), Arc::new(RecordingDelegate::default()))
.await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
Arc::new(RecordingDelegate::default()),
)
.await;
let (response_tx, response_rx) = oneshot::channel();
harness
@@ -1271,11 +1340,14 @@ async fn remote_wait_accepts_durations_longer_than_five_minutes() {
async fn queued_remote_wait_times_out_and_invalidates_the_connection() {
let mut harness = DriverHarness::start();
let session = remote_session();
harness
.open(session.clone(), Arc::new(RecordingDelegate::default()))
.await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
Arc::new(RecordingDelegate::default()),
)
.await;
let connection = Connection {
command_tx: harness.command_tx.clone(),
@@ -1317,11 +1389,14 @@ async fn queued_remote_wait_times_out_and_invalidates_the_connection() {
async fn queued_remote_termination_times_out_and_invalidates_the_connection() {
let mut harness = DriverHarness::start();
let session = remote_session();
harness
.open(session.clone(), Arc::new(RecordingDelegate::default()))
.await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
Arc::new(RecordingDelegate::default()),
)
.await;
let connection = Connection {
command_tx: harness.command_tx.clone(),
@@ -1356,11 +1431,14 @@ async fn queued_remote_termination_times_out_and_invalidates_the_connection() {
async fn cancelled_wait_is_retired_before_next_wait_is_sent() {
let mut harness = DriverHarness::start();
let session = remote_session();
harness
.open(session.clone(), Arc::new(RecordingDelegate::default()))
.await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
Arc::new(RecordingDelegate::default()),
)
.await;
let first_cancellation = CancellationToken::new();
let (first_tx, first_rx) = oneshot::channel();
@@ -1450,12 +1528,13 @@ async fn abandoned_execute_is_tracked_and_terminated_after_admission() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let cancellation = CancellationToken::new();
let (execute_tx, execute_rx) = oneshot::channel();
harness
.command_tx
.send(DriverCommand::Execute {
delegate: delegate.clone(),
session: session.clone(),
request: ExecuteRequest {
tool_call_id: "call-1".to_string(),
@@ -1547,12 +1626,13 @@ async fn delivered_but_unclaimed_execute_is_terminated_when_the_caller_is_cancel
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let cancellation = CancellationToken::new();
let (execute_tx, execute_rx) = oneshot::channel();
harness
.command_tx
.send(DriverCommand::Execute {
delegate: delegate.clone(),
session: session.clone(),
request: ExecuteRequest {
tool_call_id: "call-1".to_string(),
@@ -1635,7 +1715,13 @@ async fn delivered_but_unclaimed_execute_is_terminated_when_the_caller_is_cancel
}))
.await
.expect("cell close");
tokio::task::yield_now().await;
tokio::time::timeout(Duration::from_secs(2), async {
while Arc::strong_count(&delegate) > 1 {
tokio::task::yield_now().await;
}
})
.await
.expect("abandoned execution should release its delegate");
assert!(harness.alive.load(Ordering::Acquire));
assert_eq!(
@@ -1651,13 +1737,13 @@ async fn session_accepts_more_than_4096_cells_without_growing_a_tombstone_set()
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
for sequence in 1..=CELL_COUNT {
let request_id = i64::try_from(sequence).expect("cell sequence fits in i64") + 1;
let cell_id = sequence.to_string();
let started = harness
.start_cell(session.clone(), request_id, &cell_id)
.start_cell(session.clone(), request_id, &cell_id, delegate.clone())
.await;
harness
.event_tx
@@ -1705,11 +1791,12 @@ async fn connection_failure_closes_every_live_cell_once() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
let cleanup = harness.open(session.clone(), delegate.clone()).await;
let cleanup = harness.open(session.clone()).await;
let (execute_tx, execute_rx) = oneshot::channel();
harness
.command_tx
.send(DriverCommand::Execute {
delegate: delegate.clone(),
session,
request: ExecuteRequest {
tool_call_id: "call-1".to_string(),
@@ -1759,9 +1846,14 @@ async fn session_cleanup_does_not_wait_for_delegate_completion() {
let mut harness = DriverHarness::start();
let session = remote_session();
let (delegate, mut events_rx, release) = HeldDelegate::new();
let cleanup = harness.open(session.clone(), delegate).await;
let cleanup = harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
harness
.start_tool_delegate(&session, DelegateRequestId::new(/*value*/ 7))
@@ -1806,9 +1898,14 @@ async fn aborting_driver_marks_connection_dead_and_closes_cells() {
let mut harness = DriverHarness::start();
let session = remote_session();
let delegate = Arc::new(RecordingDelegate::default());
harness.open(session.clone(), delegate.clone()).await;
harness.open(session.clone()).await;
let _started = harness
.start_cell(session.clone(), /*request_id*/ 2, "1")
.start_cell(
session.clone(),
/*request_id*/ 2,
"1",
delegate.clone(),
)
.await;
let (wait_tx, wait_rx) = oneshot::channel();
harness
@@ -1847,9 +1944,7 @@ async fn aborting_driver_marks_connection_dead_and_closes_cells() {
async fn dropped_shutdown_waiter_does_not_abort_remote_cleanup() {
let mut harness = DriverHarness::start();
let session = remote_session();
harness
.open(session.clone(), Arc::new(RecordingDelegate::default()))
.await;
harness.open(session.clone()).await;
let (shutdown_tx, shutdown_rx) = oneshot::channel();
harness
.command_tx
@@ -1878,6 +1973,7 @@ async fn dropped_shutdown_waiter_does_not_abort_remote_cleanup() {
harness
.command_tx
.send(DriverCommand::Execute {
delegate: Arc::new(RecordingDelegate::default()),
session,
request: ExecuteRequest {
tool_call_id: "call-2".to_string(),
codex-rs/code-mode/src/remote_session_tests.rs
+11 -8
View File
@@ -62,7 +62,7 @@ async fn provider_returns_missing_host_error() {
);
let error = provider
.create_session(Arc::new(NoopCodeModeSessionDelegate))
.create_session()
.await
.err()
.expect("missing host should fail");
@@ -76,13 +76,16 @@ async fn shutdown_before_open_does_not_spawn_the_host() {
session.shutdown().await.expect("shutdown session");
let error = session
.execute(codex_code_mode_protocol::ExecuteRequest {
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: "text('unreachable')".to_string(),
yield_time_ms: None,
max_output_tokens: None,
})
.execute(
codex_code_mode_protocol::ExecuteRequest {
tool_call_id: "call-1".to_string(),
enabled_tools: Vec::new(),
source: "text('unreachable')".to_string(),
yield_time_ms: None,
max_output_tokens: None,
},
Arc::new(NoopCodeModeSessionDelegate),
)
.await
.err()
.expect("shutdown session should reject execution");