Skip to content

[codex] trace and reduce remote exec-server latency #30266

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
richardopenai wants to merge 1 commit into from
Closed

[codex] trace and reduce remote exec-server latency #30266

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
97 changes: 69 additions & 28 deletions codex-rs/core/src/unified_exec/process.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@ use tokio_util::sync::CancellationToken;

use crate::exec::is_likely_sandbox_denied;
use codex_exec_server::ExecProcess;
use codex_exec_server::ExecProcessEvent;
use codex_exec_server::ProcessSignal as ExecServerProcessSignal;
use codex_exec_server::ReadResponse as ExecReadResponse;
use codex_exec_server::StartedExecProcess;
Expand Down Expand Up @@ -425,15 +426,70 @@ impl UnifiedExecProcess {
cancellation_token,
} = output_handles;
let process = started.process;
let mut wake_rx = process.subscribe_wake();
let mut events = process.subscribe_events();
tokio::spawn(async move {
let mut after_seq = None;
let mut last_seq = 0;
loop {
match process
.read(after_seq, /*max_bytes*/ None, /*wait_ms*/ Some(0))
.await
{
Ok(response) => {
match events.recv().await {
Ok(ExecProcessEvent::Output(chunk)) => {
if chunk.seq > last_seq {
last_seq = chunk.seq;
let bytes = chunk.chunk.into_inner();
let mut guard = output_buffer.lock().await;
guard.push_chunk(bytes.clone());
drop(guard);
let _ = output_tx.send(bytes);
output_notify.notify_waiters();
}
}
Ok(ExecProcessEvent::Exited { seq, exit_code }) => {
if seq > last_seq {
last_seq = seq;
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.exited(Some(exit_code)));
}
}
Ok(ExecProcessEvent::Closed {
seq: _,
sandbox_denied,
}) => {
if sandbox_denied {
let mut state = state_tx.borrow().clone();
state.sandbox_denied = true;
let _ = state_tx.send_replace(state);
}
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
break;
}
Ok(ExecProcessEvent::Failed(message)) => {
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.failed(message));
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
break;
}
Err(broadcast::error::RecvError::Lagged(_)) => {
let response = match process
.read(
Some(last_seq),
/*max_bytes*/ None,
/*wait_ms*/ Some(0),
)
.await
{
Ok(response) => response,
Err(err) => {
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.failed(err.to_string()));
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
break;
}
};
let ExecReadResponse {
chunks,
next_seq,
Expand All @@ -443,16 +499,15 @@ impl UnifiedExecProcess {
failure,
sandbox_denied,
} = response;

for chunk in chunks {
for chunk in chunks.into_iter().filter(|chunk| chunk.seq > last_seq) {
let bytes = chunk.chunk.into_inner();
let mut guard = output_buffer.lock().await;
guard.push_chunk(bytes.clone());
drop(guard);
let _ = output_tx.send(bytes);
output_notify.notify_waiters();
}

last_seq = last_seq.max(next_seq.saturating_sub(1));
if let Some(message) = failure {
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.failed(message));
Expand All @@ -461,7 +516,6 @@ impl UnifiedExecProcess {
cancellation_token.cancel();
break;
}

if sandbox_denied {
let mut state = state_tx.borrow().clone();
state.sandbox_denied = true;
Expand All @@ -472,37 +526,24 @@ impl UnifiedExecProcess {
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.exited(exit_code));
}

if closed {
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
}

after_seq = next_seq.checked_sub(1);
if output_closed.load(Ordering::Acquire) {
break;
}
}
Err(err) => {
Err(broadcast::error::RecvError::Closed) => {
let state = state_tx.borrow().clone();
let _ = state_tx.send_replace(state.failed(err.to_string()));
let _ = state_tx.send_replace(
state.failed("exec-server process event stream closed".to_string()),
);
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
break;
}
}

if wake_rx.changed().await.is_err() {
let state = state_tx.borrow().clone();
let _ = state_tx
.send_replace(state.failed("exec-server wake channel closed".to_string()));
output_closed.store(true, Ordering::Release);
output_closed_notify.notify_waiters();
cancellation_token.cancel();
break;
}
}
})
}
Expand Down
65 changes: 48 additions & 17 deletions codex-rs/core/src/unified_exec/process_tests.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use super::process::UnifiedExecProcess;
use crate::unified_exec::UnifiedExecError;
use codex_exec_server::ExecProcess;
use codex_exec_server::ExecProcessEvent;
use codex_exec_server::ExecProcessEventReceiver;
use codex_exec_server::ExecProcessFuture;
use codex_exec_server::ExecServerError;
Expand All @@ -15,14 +16,14 @@ use std::collections::VecDeque;
use std::sync::Arc;
use tokio::sync::Mutex;
use tokio::sync::watch;
use tokio::time::Duration;

struct MockExecProcess {
process_id: ProcessId,
write_response: WriteResponse,
read_responses: Mutex<VecDeque<ReadResponse>>,
terminate_error: Option<String>,
wake_tx: watch::Sender<u64>,
events: Vec<ExecProcessEvent>,
}

impl MockExecProcess {
Expand Down Expand Up @@ -61,7 +62,7 @@ impl ExecProcess for MockExecProcess {
}

fn subscribe_events(&self) -> ExecProcessEventReceiver {
ExecProcessEventReceiver::empty()
ExecProcessEventReceiver::from_events(self.events.clone())
}

fn read(
Expand Down Expand Up @@ -100,6 +101,7 @@ async fn remote_process(
read_responses: Mutex::new(VecDeque::new()),
terminate_error,
wake_tx,
events: Vec::new(),
}),
};

Expand Down Expand Up @@ -183,29 +185,58 @@ async fn remote_process_waits_for_early_exit_event() {
write_response: WriteResponse {
status: WriteStatus::Accepted,
},
read_responses: Mutex::new(VecDeque::from([ReadResponse {
chunks: Vec::new(),
next_seq: 2,
exited: true,
exit_code: Some(17),
closed: true,
failure: None,
sandbox_denied: false,
}])),
read_responses: Mutex::new(VecDeque::new()),
terminate_error: None,
wake_tx: wake_tx.clone(),
wake_tx,
events: vec![
ExecProcessEvent::Exited {
seq: 1,
exit_code: 17,
},
ExecProcessEvent::Closed {
seq: 2,
sandbox_denied: false,
},
],
}),
};

tokio::spawn(async move {
tokio::time::sleep(Duration::from_millis(10)).await;
let _ = wake_tx.send(1);
});

let process = UnifiedExecProcess::from_exec_server_started(started)
.await
.expect("remote process should observe early exit");

assert!(process.has_exited());
assert_eq!(process.exit_code(), Some(17));
}

#[tokio::test]
async fn remote_process_preserves_streamed_sandbox_denial() {
let (wake_tx, _wake_rx) = watch::channel(0);
let started = StartedExecProcess {
process: Arc::new(MockExecProcess {
process_id: "sandbox-denied".to_string().into(),
write_response: WriteResponse {
status: WriteStatus::Accepted,
},
read_responses: Mutex::new(VecDeque::new()),
terminate_error: None,
wake_tx,
events: vec![
ExecProcessEvent::Exited {
seq: 1,
exit_code: 1,
},
ExecProcessEvent::Closed {
seq: 2,
sandbox_denied: true,
},
],
}),
};

let error = UnifiedExecProcess::from_exec_server_started(started)
.await
.expect_err("sandbox denial should be preserved");

assert!(matches!(error, UnifiedExecError::SandboxDenied { .. }));
}
14 changes: 14 additions & 0 deletions codex-rs/exec-server-protocol/src/protocol.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -534,6 +534,8 @@ pub struct ExecExitedNotification {
pub struct ExecClosedNotification {
pub process_id: ProcessId,
pub seq: u64,
#[serde(default)]
pub sandbox_denied: bool,
}

mod base64_bytes {
Expand Down Expand Up @@ -564,6 +566,7 @@ mod base64_bytes {
#[cfg(test)]
mod tests {
use super::EnvironmentInfo;
use super::ExecClosedNotification;
use super::ExecParams;
use super::FsReadFileParams;
use super::HttpRequestParams;
Expand Down Expand Up @@ -708,4 +711,15 @@ mod tests {
("req-explicit-timeout", Some(1234))
);
}

#[test]
fn closed_notification_accepts_legacy_payload_without_sandbox_denied() {
let notification: ExecClosedNotification = serde_json::from_value(serde_json::json!({
"processId": "proc-1",
"seq": 3,
}))
.expect("legacy closed notification should deserialize");

assert!(!notification.sandbox_denied);
}
}
29 changes: 27 additions & 2 deletions codex-rs/exec-server/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -334,10 +334,35 @@ Params:

```json
{
"processId": "proc-1"
"processId": "proc-1",
"seq": 3,
"sandboxDenied": false
}
```

`sandboxDenied` lets streaming clients preserve executor-side sandbox denial
detection without issuing a final `process/read` request. Older servers omit
the field, which clients interpret as `false`.

## Remote latency benchmark

Run the repeatable benchmark against an already registered remote exec-server:

```bash
CODEX_EXEC_SERVER_NOISE_REGISTRY_URL="$REGISTRY_URL" \
CODEX_EXEC_SERVER_NOISE_ENVIRONMENT_ID="$ENVIRONMENT_ID" \
CODEX_EXEC_SERVER_NOISE_AUTH_TOKEN="$AUTH_TOKEN" \
cargo run -p codex-exec-server --example remote_latency
```

It performs 5 warmups and 30 measured filesystem and process operations by
default, then prints connection timing, p50/p95 samples, per-RPC phase timing,
and bounded Rendezvous route dimensions. Set
`EXEC_SERVER_LATENCY_WARMUP_ITERATIONS` or `EXEC_SERVER_LATENCY_ITERATIONS` to
change sample counts. Process completion uses pushed events; set
`EXEC_SERVER_LATENCY_PROCESS_COMPLETION=read` to measure the legacy extra
`process/read` round trip as a control.

## Filesystem RPCs

Filesystem methods require valid `file:` URI strings and return JSON-RPC errors
Expand Down Expand Up @@ -432,5 +457,5 @@ Terminate it:
{"id":4,"method":"process/terminate","params":{"processId":"proc-1"}}
{"id":4,"result":{"running":true}}
{"method":"process/exited","params":{"processId":"proc-1","seq":3,"exitCode":0}}
{"method":"process/closed","params":{"processId":"proc-1"}}
{"method":"process/closed","params":{"processId":"proc-1","seq":4,"sandboxDenied":false}}
```
Loading
Loading