Files
codex/codex-rs/exec-server/tests/process.rs
Adam Perry @ OpenAI eeae88d8a6 Add opt-in concurrent exec-server request dispatch (#36987)
## Why

Sequential dispatch lets a long-running request block unrelated health checks
and cleanup on the same connection.

## What changed

- Add `--concurrent-requests <COUNT>` for local and remote exec-server
  connections, while retaining sequential dispatch when the option is omitted
  or set to `1`.
- Preserve handshake ordering before enabling concurrent dispatch.
- Reserve separate capacity for status, signal, terminate, and close requests so
  they remain responsive when ordinary request capacity is saturated.
- Drain queued client responses during disconnect and cancel outstanding
  request tasks during connection shutdown.

## Testing

- Cover CLI parsing and concurrency-limit validation.
- Verify default sequential behavior, pipelined handshake ordering, concurrent
  request progress, control-request responsiveness, and disconnect handling.

GitOrigin-RevId: 48e4b092e318204ed635543f01f9ee0e7df095fc
2026-08-04 22:28:15 +00:00

817 lines
26 KiB
Rust

mod common;
use std::collections::HashMap;
use codex_exec_server::EnvironmentInfo;
use codex_exec_server::EnvironmentStatus;
use codex_exec_server::EnvironmentStatusKind;
use codex_exec_server::ExecResponse;
use codex_exec_server::InitializeParams;
use codex_exec_server::InitializeResponse;
use codex_exec_server::ProcessId;
use codex_exec_server::ReadResponse;
use codex_exec_server::TerminateResponse;
use codex_exec_server::WriteResponse;
use codex_exec_server::WriteStatus;
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCResponse;
use codex_exec_server_protocol::ProcessSandboxType;
use codex_utils_path_uri::PathUri;
use common::exec_server::exec_server;
use common::exec_server::exec_server_with_env;
use pretty_assertions::assert_eq;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_starts_process_over_websocket() -> anyhow::Result<()> {
let mut server = exec_server().await?;
let process_argv = if cfg!(windows) {
vec!["cmd.exe", "/D", "/C", "exit 0"]
} else {
vec!["true"]
};
let initialize_id = server
.send_request(
"initialize",
serde_json::to_value(InitializeParams {
client_name: "exec-server-test".to_string(),
resume_session_id: None,
})?,
)
.await?;
let _ = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
)
})
.await?;
server
.send_notification("initialized", serde_json::json!({}))
.await?;
let process_start_id = server
.send_request(
"process/start",
serde_json::json!({
"processId": "proc-1",
"argv": process_argv,
"cwd": PathUri::from_host_native_path(std::env::current_dir()?)?,
"env": {},
"tty": false,
"pipeStdin": false,
"arg0": null
}),
)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_start_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
panic!("expected process/start response");
};
assert_eq!(id, process_start_id);
let process_start_response: ExecResponse = serde_json::from_value(result)?;
assert_eq!(
process_start_response,
ExecResponse {
process_id: ProcessId::from("proc-1"),
sandbox_type: Some(ProcessSandboxType::None),
}
);
server.shutdown().await?;
Ok(())
}
/// Ordinary requests run one at a time when concurrent processing is not enabled.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_runs_ordinary_requests_serially_by_default() -> anyhow::Result<()> {
let mut server = exec_server().await?;
let process_argv = if cfg!(windows) {
vec!["cmd.exe", "/D", "/C", "ping -n 601 127.0.0.1 >NUL"]
} else {
vec![
"/bin/sh",
"-c",
"parent=$PPID; while kill -0 \"$parent\" 2>/dev/null; do sleep 1; done",
]
};
let process_env = if cfg!(windows) {
serde_json::json!({ "PATH": std::env::var("PATH")? })
} else {
serde_json::json!({})
};
let initialize_id = server
.send_request(
"initialize",
serde_json::to_value(InitializeParams {
client_name: "exec-server-test".to_string(),
resume_session_id: None,
})?,
)
.await?;
let _ = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
)
})
.await?;
server
.send_notification("initialized", serde_json::json!({}))
.await?;
let process_start_id = server
.send_request(
"process/start",
serde_json::json!({
"processId": "proc-serial-read",
"argv": process_argv,
"cwd": PathUri::from_host_native_path(std::env::current_dir()?)?,
"env": process_env,
"tty": false,
"pipeStdin": false,
"arg0": null
}),
)
.await?;
let _ = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_start_id
)
})
.await?;
let read_id = server
.send_request(
"process/read",
serde_json::json!({
"processId": "proc-serial-read",
"afterSeq": null,
"maxBytes": null,
"waitMs": 250
}),
)
.await?;
let queued_environment_info_id = server
.send_request("environment/info", serde_json::json!({}))
.await?;
let queued_start_id = server
.send_request(
"process/start",
serde_json::json!({
"processId": "proc-serial-queued",
"argv": process_argv,
"cwd": PathUri::from_host_native_path(std::env::current_dir()?)?,
"env": process_env,
"tty": false,
"pipeStdin": false,
"arg0": null
}),
)
.await?;
let response = server
.wait_for_event(|event| matches!(event, JSONRPCMessage::Response(_)))
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { id, .. }) = response else {
panic!("expected the blocked process/read to finish before the queued process/start");
};
assert_eq!(id, read_id);
let response = server
.wait_for_event(|event| matches!(event, JSONRPCMessage::Response(_)))
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
panic!("expected the queued environment/info response after process/read");
};
assert_eq!(id, queued_environment_info_id);
let _: EnvironmentInfo = serde_json::from_value(result)?;
let response = server
.wait_for_event(|event| matches!(event, JSONRPCMessage::Response(_)))
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
panic!("expected the queued process/start response after process/read");
};
assert_eq!(id, queued_start_id);
assert_eq!(
serde_json::from_value::<ExecResponse>(result)?,
ExecResponse {
process_id: ProcessId::from("proc-serial-queued"),
sandbox_type: Some(ProcessSandboxType::None),
}
);
for process_id in ["proc-serial-read", "proc-serial-queued"] {
let terminate_id = server
.send_request(
"process/terminate",
serde_json::json!({ "processId": process_id }),
)
.await?;
server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &terminate_id
)
})
.await?;
}
server.shutdown().await?;
Ok(())
}
/// A long read cannot block health checks; saturated ordinary work cannot block health checks or cleanup.
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_keeps_control_requests_live_during_long_reads_and_queued_requests()
-> anyhow::Result<()> {
let mut server = exec_server_with_env(
std::iter::empty::<(&str, &str)>(),
&["--concurrent-requests", "32"],
)
.await?;
let process_argv = if cfg!(windows) {
vec!["cmd.exe", "/D", "/C", "ping -n 601 127.0.0.1 >NUL"]
} else {
vec![
"/bin/sh",
"-c",
"parent=$PPID; while kill -0 \"$parent\" 2>/dev/null; do sleep 1; done",
]
};
let process_env = if cfg!(windows) {
serde_json::json!({ "PATH": std::env::var("PATH")? })
} else {
serde_json::json!({})
};
let initialize_id = server
.send_request(
"initialize",
serde_json::to_value(InitializeParams {
client_name: "exec-server-test".to_string(),
resume_session_id: None,
})?,
)
.await?;
assert!(matches!(
server.next_event().await?,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == initialize_id
));
server
.send_notification("initialized", serde_json::json!({}))
.await?;
let process_start_id = server
.send_request(
"process/start",
serde_json::json!({
"processId": "proc-capacity",
"argv": process_argv,
"cwd": PathUri::from_host_native_path(std::env::current_dir()?)?,
"env": process_env,
"tty": false,
"pipeStdin": false,
"arg0": null
}),
)
.await?;
assert!(matches!(
server.next_event().await?,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == process_start_id
));
let read_params = serde_json::json!({
"processId": "proc-capacity",
"afterSeq": null,
"maxBytes": null,
"waitMs": 600_000
});
server
.send_request("process/read", read_params.clone())
.await?;
let concurrent_start_id = server
.send_request(
"process/start",
serde_json::json!({
"processId": "proc-concurrent",
"argv": process_argv,
"cwd": PathUri::from_host_native_path(std::env::current_dir()?)?,
"env": process_env,
"tty": false,
"pipeStdin": false,
"arg0": null
}),
)
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = server.next_event().await?
else {
panic!("expected process/start to finish before the pending process/read");
};
assert_eq!(id, concurrent_start_id);
assert_eq!(
serde_json::from_value::<ExecResponse>(result)?,
ExecResponse {
process_id: ProcessId::from("proc-concurrent"),
sandbox_type: Some(ProcessSandboxType::None),
}
);
for _ in 1..32 {
server
.send_request("process/read", read_params.clone())
.await?;
}
let queued_read_id = server.send_request("process/read", read_params).await?;
let environment_info_id = server
.send_request("environment/info", serde_json::json!({}))
.await?;
let response = server
.wait_for_event(|event| matches!(event, JSONRPCMessage::Response(_)))
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
panic!("expected environment/info response at regular request capacity");
};
assert_eq!(id, environment_info_id);
let _: EnvironmentInfo = serde_json::from_value(result)?;
let environment_status_id = server
.send_request("environment/status", serde_json::json!({}))
.await?;
let response = server
.wait_for_event(|event| matches!(event, JSONRPCMessage::Response(_)))
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
panic!("expected environment/status response at regular request capacity");
};
assert_eq!(id, environment_status_id);
assert_eq!(
serde_json::from_value::<EnvironmentStatus>(result)?,
EnvironmentStatus {
status: EnvironmentStatusKind::Ready,
}
);
let terminate_id = server
.send_request(
"process/terminate",
serde_json::json!({ "processId": "proc-capacity" }),
)
.await?;
let mut terminate_response = None;
let mut queued_read_completed = false;
while terminate_response.is_none() || !queued_read_completed {
match server.next_event().await? {
JSONRPCMessage::Response(JSONRPCResponse { id, result }) if id == terminate_id => {
terminate_response = Some(serde_json::from_value::<TerminateResponse>(result)?);
}
JSONRPCMessage::Response(JSONRPCResponse { id, result }) if id == queued_read_id => {
let _: ReadResponse = serde_json::from_value(result)?;
queued_read_completed = true;
}
JSONRPCMessage::Error(error) => {
anyhow::bail!("unexpected error while waiting for queued requests: {error:?}");
}
JSONRPCMessage::Request(_)
| JSONRPCMessage::Response(_)
| JSONRPCMessage::Notification(_) => {}
}
}
assert_eq!(
terminate_response,
Some(TerminateResponse { running: true })
);
let terminate_id = server
.send_request(
"process/terminate",
serde_json::json!({ "processId": "proc-concurrent" }),
)
.await?;
server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &terminate_id
)
})
.await?;
server.shutdown().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_defaults_omitted_pipe_stdin_to_closed_stdin() -> anyhow::Result<()> {
let mut server = exec_server().await?;
let process_argv = if cfg!(windows) {
vec!["cmd.exe", "/D", "/C", "ping -n 2 127.0.0.1 >NUL"]
} else {
vec![
"/bin/sh",
"-c",
"sleep 0.3; if IFS= read -r line; then printf 'read:%s\\n' \"$line\"; else printf 'eof\\n'; fi",
]
};
let process_env = if cfg!(windows) {
serde_json::json!({ "PATH": std::env::var("PATH")? })
} else {
serde_json::json!({})
};
let initialize_id = server
.send_request(
"initialize",
serde_json::to_value(InitializeParams {
client_name: "exec-server-test".to_string(),
resume_session_id: None,
})?,
)
.await?;
let _ = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
)
})
.await?;
server
.send_notification("initialized", serde_json::json!({}))
.await?;
let process_start_id = server
.send_request(
"process/start",
serde_json::json!({
"processId": "proc-default-stdin",
"argv": process_argv,
"cwd": PathUri::from_host_native_path(std::env::current_dir()?)?,
"env": process_env,
"tty": false,
"arg0": null
}),
)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_start_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { result, .. }) = response else {
panic!("expected process/start response");
};
let process_start_response: ExecResponse = serde_json::from_value(result)?;
assert_eq!(
process_start_response,
ExecResponse {
process_id: ProcessId::from("proc-default-stdin"),
sandbox_type: Some(ProcessSandboxType::None),
}
);
let write_id = server
.send_request(
"process/write",
serde_json::json!({
"processId": "proc-default-stdin",
"chunk": "aWdub3JlZAo=",
"writeId": "write-default-stdin"
}),
)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &write_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { result, .. }) = response else {
panic!("expected process/write response");
};
let write_response: WriteResponse = serde_json::from_value(result)?;
assert_eq!(
write_response,
WriteResponse {
status: WriteStatus::StdinClosed
}
);
server.shutdown().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_dedupes_retried_process_write_ids() -> anyhow::Result<()> {
let mut server = exec_server().await?;
let process_argv = if cfg!(windows) {
vec![
"powershell.exe",
"-NoProfile",
"-NonInteractive",
"-Command",
"[Console]::Out.WriteLine('line:' + [Console]::In.ReadLine()); [Console]::Out.WriteLine('line:' + [Console]::In.ReadLine())",
]
} else {
vec![
"/bin/sh",
"-c",
"IFS= read -r first; printf 'line:%s\\n' \"$first\"; IFS= read -r second; printf 'line:%s\\n' \"$second\"",
]
};
let process_env = if cfg!(windows) {
serde_json::to_value(std::env::vars().collect::<HashMap<_, _>>())?
} else {
serde_json::json!({})
};
let initialize_id = server
.send_request(
"initialize",
serde_json::to_value(InitializeParams {
client_name: "exec-server-test".to_string(),
resume_session_id: None,
})?,
)
.await?;
let _ = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
)
})
.await?;
server
.send_notification("initialized", serde_json::json!({}))
.await?;
let process_start_id = server
.send_request(
"process/start",
serde_json::json!({
"processId": "proc-write-id",
"argv": process_argv,
"cwd": PathUri::from_host_native_path(std::env::current_dir()?)?,
"env": process_env,
"tty": false,
"pipeStdin": true,
"arg0": null
}),
)
.await?;
let _ = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_start_id
)
})
.await?;
for (write_id, chunk) in [
("write-1", "Zmlyc3QK"),
("write-1", "Zmlyc3QK"),
("write-2", "c2Vjb25kCg=="),
] {
let request_id = server
.send_request(
"process/write",
serde_json::json!({
"processId": "proc-write-id",
"chunk": chunk,
"writeId": write_id
}),
)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &request_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { result, .. }) = response else {
panic!("expected process/write response");
};
let write_response: WriteResponse = serde_json::from_value(result)?;
assert_eq!(
write_response,
WriteResponse {
status: WriteStatus::Accepted
}
);
}
let mut after_seq = None;
let mut output = Vec::new();
for _ in 0..5 {
let read_id = server
.send_request(
"process/read",
serde_json::json!({
"processId": "proc-write-id",
"afterSeq": after_seq,
"maxBytes": null,
"waitMs": 1000
}),
)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &read_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { result, .. }) = response else {
panic!("expected process/read response");
};
let read_response: ReadResponse = serde_json::from_value(result)?;
output.extend(
read_response
.chunks
.into_iter()
.flat_map(|chunk| chunk.chunk.into_inner()),
);
after_seq = Some(read_response.next_seq.saturating_sub(1));
if read_response.closed
|| output.ends_with(b"line:second\n")
|| output.ends_with(b"line:second\r\n")
{
break;
}
}
assert_eq!(
String::from_utf8(output)?.replace("\r\n", "\n"),
"line:first\nline:second\n".to_string()
);
server.shutdown().await?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_resumes_detached_session_without_killing_processes() -> anyhow::Result<()> {
let mut server = exec_server().await?;
let process_argv = if cfg!(windows) {
vec!["cmd.exe", "/D", "/C", "ping -n 6 127.0.0.1 >NUL"]
} else {
vec!["/bin/sh", "-c", "sleep 5"]
};
let process_env = if cfg!(windows) {
serde_json::json!({ "PATH": std::env::var("PATH")? })
} else {
serde_json::json!({})
};
let initialize_id = server
.send_request(
"initialize",
serde_json::to_value(InitializeParams {
client_name: "exec-server-test".to_string(),
resume_session_id: None,
})?,
)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &initialize_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { result, .. }) = response else {
panic!("expected initialize response");
};
let initialize_response: InitializeResponse = serde_json::from_value(result)?;
server
.send_notification("initialized", serde_json::json!({}))
.await?;
let process_start_id = server
.send_request(
"process/start",
serde_json::json!({
"processId": "proc-resume",
"argv": process_argv,
"cwd": PathUri::from_host_native_path(std::env::current_dir()?)?,
"env": process_env,
"tty": false,
"pipeStdin": false,
"arg0": null
}),
)
.await?;
let _ = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_start_id
)
})
.await?;
server.disconnect_websocket().await?;
server.reconnect_websocket().await?;
let resume_initialize_id = server
.send_request(
"initialize",
serde_json::to_value(InitializeParams {
client_name: "exec-server-test".to_string(),
resume_session_id: Some(initialize_response.session_id.clone()),
})?,
)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &resume_initialize_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { result, .. }) = response else {
panic!("expected resume initialize response");
};
let resumed_response: InitializeResponse = serde_json::from_value(result)?;
assert_eq!(resumed_response, initialize_response);
server
.send_notification("initialized", serde_json::json!({}))
.await?;
let process_read_id = server
.send_request(
"process/read",
serde_json::json!({
"processId": "proc-resume",
"afterSeq": null,
"maxBytes": null,
"waitMs": 0
}),
)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &process_read_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { result, .. }) = response else {
panic!("expected process/read response");
};
let process_read_response: ReadResponse = serde_json::from_value(result)?;
assert!(process_read_response.failure.is_none());
assert!(!process_read_response.exited);
assert!(!process_read_response.closed);
let terminate_id = server
.send_request(
"process/terminate",
serde_json::json!({
"processId": "proc-resume"
}),
)
.await?;
let response = server
.wait_for_event(|event| {
matches!(
event,
JSONRPCMessage::Response(JSONRPCResponse { id, .. }) if id == &terminate_id
)
})
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { result, .. }) = response else {
panic!("expected process/terminate response");
};
let terminate_response: TerminateResponse = serde_json::from_value(result)?;
assert_eq!(terminate_response, TerminateResponse { running: true });
server.shutdown().await?;
Ok(())
}