clone/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-04-24 06:35:50 +00:00
Code Issues Packages Projects Releases Wiki Activity

feat: exec-server prep for unified exec

Browse Source
This commit is contained in:
jif-oai
2026-03-24 22:14:43 +00:00
parent 621862a7d1
commit 3ce6a55209
11 changed files with 528 additions and 155 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
codex-rs/Cargo.lock generated
View File

@@ -2042,6 +2042,7 @@ dependencies = [
"tokio",
"tokio-tungstenite",
"tracing",
"tracing-subscriber",
]
[[package]]

1
codex-rs/exec-server/Cargo.toml
View File

@@ -38,6 +38,7 @@ tokio = { workspace = true, features = [
] }
tokio-tungstenite = { workspace = true }
tracing = { workspace = true }
tracing-subscriber = { workspace = true, features = ["env-filter", "fmt"] }
[dev-dependencies]
anyhow = { workspace = true }

7
codex-rs/exec-server/src/bin/codex-exec-server.rs
View File

@@ -1,4 +1,5 @@
use clap::Parser;
use tracing_subscriber::EnvFilter;
#[derive(Debug, Parser)]
struct ExecServerArgs {
@@ -13,6 +14,12 @@ struct ExecServerArgs {
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let env_filter = EnvFilter::try_from_default_env().unwrap_or_else(|_| EnvFilter::new("warn"));
let _ = tracing_subscriber::fmt()
.with_env_filter(env_filter)
.with_target(false)
.try_init();
let args = ExecServerArgs::parse();
codex_exec_server::run_main_with_listen_url(&args.listen).await
}

58
codex-rs/exec-server/src/client.rs
View File

@@ -1,3 +1,4 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
@@ -17,6 +18,7 @@ use codex_app_server_protocol::FsWriteFileParams;
use codex_app_server_protocol::FsWriteFileResponse;
use codex_app_server_protocol::JSONRPCNotification;
use serde_json::Value;
use tokio::sync::Mutex;
use tokio::sync::broadcast;
use tokio::time::timeout;
use tokio_tungstenite::connect_async;
@@ -26,13 +28,15 @@ use tracing::warn;
use crate::client_api::ExecServerClientConnectOptions;
use crate::client_api::RemoteExecServerConnectArgs;
use crate::connection::JsonRpcConnection;
use crate::process::ExecServerEvent;
use crate::process::ExecSessionEvent;
use crate::protocol::EXEC_CLOSED_METHOD;
use crate::protocol::EXEC_EXITED_METHOD;
use crate::protocol::EXEC_METHOD;
use crate::protocol::EXEC_OUTPUT_DELTA_METHOD;
use crate::protocol::EXEC_READ_METHOD;
use crate::protocol::EXEC_TERMINATE_METHOD;
use crate::protocol::EXEC_WRITE_METHOD;
use crate::protocol::ExecClosedNotification;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecParams;
@@ -92,7 +96,7 @@ impl RemoteExecServerConnectArgs {
struct Inner {
client: RpcClient,
events_tx: broadcast::Sender<ExecServerEvent>,
sessions: Mutex<HashMap<String, broadcast::Sender<ExecSessionEvent>>>,
reader_task: tokio::task::JoinHandle<()>,
}
@@ -158,10 +162,6 @@ impl ExecServerClient {
.await
}
pub fn event_receiver(&self) -> broadcast::Receiver<ExecServerEvent> {
self.inner.events_tx.subscribe()
}
pub async fn initialize(
&self,
options: ExecServerClientConnectOptions,
@@ -307,6 +307,25 @@ impl ExecServerClient {
.map_err(Into::into)
}
pub(crate) async fn register_session(
&self,
process_id: &str,
) -> Result<broadcast::Receiver<ExecSessionEvent>, ExecServerError> {
let (events_tx, events_rx) = broadcast::channel(256);
let mut sessions = self.inner.sessions.lock().await;
if sessions.contains_key(process_id) {
return Err(ExecServerError::Protocol(format!(
"session already registered for process {process_id}"
)));
}
sessions.insert(process_id.to_string(), events_tx);
Ok(events_rx)
}
pub(crate) async fn unregister_session(&self, process_id: &str) {
self.inner.sessions.lock().await.remove(process_id);
}
async fn connect(
connection: JsonRpcConnection,
options: ExecServerClientConnectOptions,
@@ -338,7 +357,7 @@ impl ExecServerClient {
Inner {
client: rpc_client,
events_tx: broadcast::channel(256).0,
sessions: Mutex::new(HashMap::new()),
reader_task,
}
});
@@ -378,12 +397,33 @@ async fn handle_server_notification(
EXEC_OUTPUT_DELTA_METHOD => {
let params: ExecOutputDeltaNotification =
serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
let _ = inner.events_tx.send(ExecServerEvent::OutputDelta(params));
let events_tx = { inner.sessions.lock().await.get(&params.process_id).cloned() };
if let Some(events_tx) = events_tx {
let _ = events_tx.send(ExecSessionEvent::Output {
seq: params.seq,
stream: params.stream,
chunk: params.chunk.into_inner(),
});
}
}
EXEC_EXITED_METHOD => {
let params: ExecExitedNotification =
serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
let _ = inner.events_tx.send(ExecServerEvent::Exited(params));
let events_tx = { inner.sessions.lock().await.get(&params.process_id).cloned() };
if let Some(events_tx) = events_tx {
let _ = events_tx.send(ExecSessionEvent::Exited {
seq: params.seq,
exit_code: params.exit_code,
});
}
}
EXEC_CLOSED_METHOD => {
let params: ExecClosedNotification =
serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
let events_tx = { inner.sessions.lock().await.remove(&params.process_id) };
if let Some(events_tx) = events_tx {
let _ = events_tx.send(ExecSessionEvent::Closed { seq: params.seq });
}
}
other => {
debug!("ignoring unknown exec-server notification: {other}");

53
codex-rs/exec-server/src/environment.rs
View File

@@ -6,19 +6,19 @@ use crate::RemoteExecServerConnectArgs;
use crate::file_system::ExecutorFileSystem;
use crate::local_file_system::LocalFileSystem;
use crate::local_process::LocalProcess;
use crate::process::ExecProcess;
use crate::process::ExecBackend;
use crate::remote_file_system::RemoteFileSystem;
use crate::remote_process::RemoteProcess;
pub trait ExecutorEnvironment: Send + Sync {
fn get_executor(&self) -> Arc<dyn ExecProcess>;
fn get_exec_backend(&self) -> Arc<dyn ExecBackend>;
}
#[derive(Clone)]
pub struct Environment {
experimental_exec_server_url: Option<String>,
remote_exec_server_client: Option<ExecServerClient>,
executor: Arc<dyn ExecProcess>,
exec_backend: Arc<dyn ExecBackend>,
}
impl Default for Environment {
@@ -34,7 +34,7 @@ impl Default for Environment {
Self {
experimental_exec_server_url: None,
remote_exec_server_client: None,
executor: Arc::new(local_process),
exec_backend: Arc::new(local_process),
}
}
}
@@ -68,24 +68,24 @@ impl Environment {
None
};
let executor: Arc<dyn ExecProcess> = if let Some(client) = remote_exec_server_client.clone()
{
Arc::new(RemoteProcess::new(client))
} else {
let local_process = LocalProcess::default();
local_process
.initialize()
.map_err(|err| ExecServerError::Protocol(err.message))?;
local_process
.initialized()
.map_err(ExecServerError::Protocol)?;
Arc::new(local_process)
};
let exec_backend: Arc<dyn ExecBackend> =
if let Some(client) = remote_exec_server_client.clone() {
Arc::new(RemoteProcess::new(client))
} else {
let local_process = LocalProcess::default();
local_process
.initialize()
.map_err(|err| ExecServerError::Protocol(err.message))?;
local_process
.initialized()
.map_err(ExecServerError::Protocol)?;
Arc::new(local_process)
};
Ok(Self {
experimental_exec_server_url,
remote_exec_server_client,
executor,
exec_backend,
})
}
@@ -93,8 +93,8 @@ impl Environment {
self.experimental_exec_server_url.as_deref()
}
pub fn get_executor(&self) -> Arc<dyn ExecProcess> {
Arc::clone(&self.executor)
pub fn get_exec_backend(&self) -> Arc<dyn ExecBackend> {
Arc::clone(&self.exec_backend)
}
pub fn get_filesystem(&self) -> Arc<dyn ExecutorFileSystem> {
@@ -107,8 +107,8 @@ impl Environment {
}
impl ExecutorEnvironment for Environment {
fn get_executor(&self) -> Arc<dyn ExecProcess> {
Arc::clone(&self.executor)
fn get_exec_backend(&self) -> Arc<dyn ExecBackend> {
Arc::clone(&self.exec_backend)
}
}
@@ -130,7 +130,7 @@ mod tests {
let environment = Environment::default();
let response = environment
.get_executor()
.get_exec_backend()
.start(crate::ExecParams {
process_id: "default-env-proc".to_string(),
argv: vec!["true".to_string()],
@@ -142,11 +142,6 @@ mod tests {
.await
.expect("start process");
assert_eq!(
response,
crate::ExecResponse {
process_id: "default-env-proc".to_string(),
}
);
assert_eq!(response.process_id, "default-env-proc");
}
}

4
codex-rs/exec-server/src/lib.rs
View File

@@ -39,8 +39,10 @@ pub use file_system::FileMetadata;
pub use file_system::FileSystemResult;
pub use file_system::ReadDirectoryEntry;
pub use file_system::RemoveOptions;
pub use process::ExecBackend;
pub use process::ExecProcess;
pub use process::ExecServerEvent;
pub use process::ExecSessionEvent;
pub use protocol::ExecClosedNotification;
pub use protocol::ExecExitedNotification;
pub use protocol::ExecOutputDeltaNotification;
pub use protocol::ExecOutputStream;

281
codex-rs/exec-server/src/local_process.rs
View File

@@ -1,6 +1,7 @@
use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::time::Duration;
@@ -15,9 +16,12 @@ use tokio::sync::broadcast;
use tokio::sync::mpsc;
use tracing::warn;
use crate::ExecBackend;
use crate::ExecProcess;
use crate::ExecServerError;
use crate::ExecServerEvent;
use crate::ExecSessionEvent;
use crate::protocol::EXEC_CLOSED_METHOD;
use crate::protocol::ExecClosedNotification;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecOutputStream;
@@ -60,6 +64,9 @@ struct RunningProcess {
next_seq: u64,
exit_code: Option<i32>,
output_notify: Arc<Notify>,
session_events_tx: broadcast::Sender<ExecSessionEvent>,
open_streams: usize,
closed: bool,
}
enum ProcessEntry {
@@ -69,7 +76,6 @@ enum ProcessEntry {
struct Inner {
notifications: RpcNotificationSender,
events_tx: broadcast::Sender<ExecServerEvent>,
processes: Mutex<HashMap<String, ProcessEntry>>,
initialize_requested: AtomicBool,
initialized: AtomicBool,
@@ -80,6 +86,13 @@ pub(crate) struct LocalProcess {
inner: Arc<Inner>,
}
#[derive(Clone)]
struct LocalExecProcess {
process_id: String,
events: StdMutex<broadcast::Receiver<ExecSessionEvent>>,
backend: LocalProcess,
}
impl Default for LocalProcess {
fn default() -> Self {
let (outgoing_tx, mut outgoing_rx) =
@@ -94,7 +107,6 @@ impl LocalProcess {
Self {
inner: Arc::new(Inner {
notifications,
events_tx: broadcast::channel(EVENT_CHANNEL_CAPACITY).0,
processes: Mutex::new(HashMap::new()),
initialize_requested: AtomicBool::new(false),
initialized: AtomicBool::new(false),
@@ -113,6 +125,10 @@ impl LocalProcess {
})
.collect::<Vec<_>>()
};
warn!(
remaining_processes = remaining.len(),
"exec-server shutting down local process handler"
);
for process in remaining {
process.session.terminate();
}
@@ -124,6 +140,7 @@ impl LocalProcess {
"initialize may only be sent once per connection".to_string(),
));
}
warn!("exec-server received initialize request");
Ok(InitializeResponse {})
}
@@ -132,6 +149,7 @@ impl LocalProcess {
return Err("received `initialized` notification before `initialize`".into());
}
self.inner.initialized.store(true, Ordering::SeqCst);
warn!("exec-server received initialized notification");
Ok(())
}
@@ -152,9 +170,19 @@ impl LocalProcess {
Ok(())
}
pub(crate) async fn exec(&self, params: ExecParams) -> Result<ExecResponse, JSONRPCErrorError> {
async fn start_process(
&self,
params: ExecParams,
) -> Result<(ExecResponse, broadcast::Receiver<ExecSessionEvent>), JSONRPCErrorError> {
self.require_initialized_for("exec")?;
let process_id = params.process_id.clone();
warn!(
process_id = %process_id,
tty = params.tty,
cwd = %params.cwd.display(),
argv = ?params.argv,
"exec-server starting process"
);
let (program, args) = params
.argv
@@ -198,11 +226,17 @@ impl LocalProcess {
if matches!(process_map.get(&process_id), Some(ProcessEntry::Starting)) {
process_map.remove(&process_id);
}
warn!(
process_id = %process_id,
error = %err,
"exec-server failed to spawn process"
);
return Err(internal_error(err.to_string()));
}
};
let output_notify = Arc::new(Notify::new());
let (session_events_tx, session_events_rx) = broadcast::channel(EVENT_CHANNEL_CAPACITY);
{
let mut process_map = self.inner.processes.lock().await;
process_map.insert(
@@ -215,6 +249,9 @@ impl LocalProcess {
next_seq: 1,
exit_code: None,
output_notify: Arc::clone(&output_notify),
session_events_tx,
open_streams: 2,
closed: false,
})),
);
}
@@ -248,7 +285,18 @@ impl LocalProcess {
output_notify,
));
Ok(ExecResponse { process_id })
warn!(
process_id = %process_id,
tty = params.tty,
"exec-server started process"
);
Ok((ExecResponse { process_id }, session_events_rx))
}
pub(crate) async fn exec(&self, params: ExecParams) -> Result<ExecResponse, JSONRPCErrorError> {
self.start_process(params)
.await
.map(|(response, _)| response)
}
pub(crate) async fn exec_read(
@@ -256,6 +304,7 @@ impl LocalProcess {
params: ReadParams,
) -> Result<ReadResponse, JSONRPCErrorError> {
self.require_initialized_for("exec")?;
let process_id = params.process_id.clone();
let after_seq = params.after_seq.unwrap_or(0);
let max_bytes = params.max_bytes.unwrap_or(usize::MAX);
let wait = Duration::from_millis(params.wait_ms.unwrap_or(0));
@@ -309,6 +358,23 @@ impl LocalProcess {
|| response.exited
|| tokio::time::Instant::now() >= deadline
{
let total_bytes: usize = response
.chunks
.iter()
.map(|chunk| chunk.chunk.0.len())
.sum();
warn!(
process_id = %process_id,
after_seq,
next_seq = response.next_seq,
chunk_count = response.chunks.len(),
total_bytes,
exited = response.exited,
exit_code = ?response.exit_code,
wait_ms = params.wait_ms.unwrap_or(0),
max_bytes,
"exec-server returning process/read response"
);
return Ok(response);
}
@@ -325,6 +391,8 @@ impl LocalProcess {
params: WriteParams,
) -> Result<WriteResponse, JSONRPCErrorError> {
self.require_initialized_for("exec")?;
let process_id = params.process_id.clone();
let input_bytes = params.chunk.0.len();
let writer_tx = {
let process_map = self.inner.processes.lock().await;
let process = process_map.get(&params.process_id).ok_or_else(|| {
@@ -350,6 +418,11 @@ impl LocalProcess {
.await
.map_err(|_| internal_error("failed to write to process stdin".to_string()))?;
warn!(
process_id = %process_id,
input_bytes,
"exec-server wrote stdin to process"
);
Ok(WriteResponse { accepted: true })
}
@@ -358,6 +431,7 @@ impl LocalProcess {
params: TerminateParams,
) -> Result<TerminateResponse, JSONRPCErrorError> {
self.require_initialized_for("exec")?;
let process_id = params.process_id.clone();
let running = {
let process_map = self.inner.processes.lock().await;
match process_map.get(&params.process_id) {
@@ -372,43 +446,78 @@ impl LocalProcess {
}
};
warn!(
process_id = %process_id,
running,
"exec-server processed terminate request"
);
Ok(TerminateResponse { running })
}
}
#[async_trait]
impl ExecProcess for LocalProcess {
async fn start(&self, params: ExecParams) -> Result<ExecResponse, ExecServerError> {
self.exec(params).await.map_err(map_handler_error)
impl ExecBackend for LocalProcess {
async fn start(&self, params: ExecParams) -> Result<Arc<dyn ExecProcess>, ExecServerError> {
let (response, events) = self
.start_process(params)
.await
.map_err(map_handler_error)?;
Ok(Arc::new(LocalExecProcess {
process_id: response.process_id,
events: StdMutex::new(events),
backend: self.clone(),
}))
}
}
#[async_trait]
impl ExecProcess for LocalExecProcess {
fn process_id(&self) -> &str {
&self.process_id
}
async fn read(&self, params: ReadParams) -> Result<ReadResponse, ExecServerError> {
self.exec_read(params).await.map_err(map_handler_error)
fn subscribe(&self) -> broadcast::Receiver<ExecSessionEvent> {
self
.events
.lock()
.expect("local exec process events mutex should not be poisoned")
.resubscribe()
}
async fn write(
&self,
process_id: &str,
chunk: Vec<u8>,
) -> Result<WriteResponse, ExecServerError> {
self.exec_write(WriteParams {
process_id: process_id.to_string(),
chunk: chunk.into(),
})
.await
.map_err(map_handler_error)
async fn write_stdin(&self, chunk: Vec<u8>) -> Result<(), ExecServerError> {
self.backend.write_stdin(&self.process_id, chunk).await
}
async fn terminate(&self, process_id: &str) -> Result<TerminateResponse, ExecServerError> {
async fn terminate(&self) -> Result<(), ExecServerError> {
self.backend.terminate(&self.process_id).await
}
}
impl LocalProcess {
async fn write_stdin(&self, process_id: &str, chunk: Vec<u8>) -> Result<(), ExecServerError> {
let response = self
.exec_write(WriteParams {
process_id: process_id.to_string(),
chunk: chunk.into(),
})
.await
.map_err(map_handler_error)?;
if response.accepted {
Ok(())
} else {
Err(ExecServerError::Protocol(format!(
"exec-server did not accept stdin for process {process_id}"
)))
}
}
async fn terminate(&self, process_id: &str) -> Result<(), ExecServerError> {
self.terminate_process(TerminateParams {
process_id: process_id.to_string(),
})
.await
.map_err(map_handler_error)
}
fn subscribe_events(&self) -> broadcast::Receiver<ExecServerEvent> {
self.inner.events_tx.subscribe()
.map_err(map_handler_error)?;
Ok(())
}
}
@@ -427,6 +536,7 @@ async fn stream_output(
output_notify: Arc<Notify>,
) {
while let Some(chunk) = receiver.recv().await {
let chunk_len = chunk.len();
let notification = {
let mut processes = inner.processes.lock().await;
let Some(entry) = processes.get_mut(&process_id) else {
@@ -452,17 +562,26 @@ async fn stream_output(
"retained output cap exceeded for process {process_id}; dropping oldest output"
);
}
let event = ExecSessionEvent::Output {
seq,
stream,
chunk: chunk.clone(),
};
let _ = process.session_events_tx.send(event);
ExecOutputDeltaNotification {
process_id: process_id.clone(),
seq,
stream,
chunk: chunk.into(),
}
};
warn!(
process_id = %process_id,
?stream,
chunk_bytes = chunk_len,
"exec-server emitted output chunk"
);
output_notify.notify_waiters();
let _ = inner
.events_tx
.send(ExecServerEvent::OutputDelta(notification.clone()));
if inner
.notifications
.notify(crate::protocol::EXEC_OUTPUT_DELTA_METHOD, &notification)
@@ -472,6 +591,8 @@ async fn stream_output(
break;
}
}
finish_output_stream(process_id, inner).await;
}
async fn watch_exit(
@@ -481,29 +602,42 @@ async fn watch_exit(
output_notify: Arc<Notify>,
) {
let exit_code = exit_rx.await.unwrap_or(-1);
{
warn!(
process_id = %process_id,
exit_code,
"exec-server observed process exit"
);
let notification = {
let mut processes = inner.processes.lock().await;
if let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) {
let seq = process.next_seq;
process.next_seq += 1;
process.exit_code = Some(exit_code);
let _ = process
.session_events_tx
.send(ExecSessionEvent::Exited { seq, exit_code });
Some(ExecExitedNotification {
process_id: process_id.clone(),
seq,
exit_code,
})
} else {
None
}
}
output_notify.notify_waiters();
let notification = ExecExitedNotification {
process_id: process_id.clone(),
exit_code,
};
let _ = inner
.events_tx
.send(ExecServerEvent::Exited(notification.clone()));
if inner
.notifications
.notify(crate::protocol::EXEC_EXITED_METHOD, &notification)
.await
.is_err()
output_notify.notify_waiters();
if let Some(notification) = notification
&& inner
.notifications
.notify(crate::protocol::EXEC_EXITED_METHOD, &notification)
.await
.is_err()
{
return;
}
maybe_emit_closed(process_id.clone(), Arc::clone(&inner)).await;
tokio::time::sleep(EXITED_PROCESS_RETENTION).await;
let mut processes = inner.processes.lock().await;
if matches!(
@@ -511,5 +645,62 @@ async fn watch_exit(
Some(ProcessEntry::Running(process)) if process.exit_code == Some(exit_code)
) {
processes.remove(&process_id);
warn!(
process_id = %process_id,
exit_code,
"exec-server evicted exited process from retention cache"
);
}
}
async fn finish_output_stream(process_id: String, inner: Arc<Inner>) {
{
let mut processes = inner.processes.lock().await;
let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) else {
return;
};
if process.open_streams > 0 {
process.open_streams -= 1;
}
}
maybe_emit_closed(process_id, inner).await;
}
async fn maybe_emit_closed(process_id: String, inner: Arc<Inner>) {
let notification = {
let mut processes = inner.processes.lock().await;
let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) else {
return;
};
if process.closed || process.open_streams != 0 || process.exit_code.is_none() {
return;
}
process.closed = true;
let seq = process.next_seq;
process.next_seq += 1;
let _ = process
.session_events_tx
.send(ExecSessionEvent::Closed { seq });
Some(ExecClosedNotification {
process_id: process_id.clone(),
seq,
})
};
let Some(notification) = notification else {
return;
};
if inner
.notifications
.notify(EXEC_CLOSED_METHOD, &notification)
.await
.is_err()
{
return;
}
}

45
codex-rs/exec-server/src/process.rs
View File

@@ -1,35 +1,40 @@
use std::sync::Arc;
use async_trait::async_trait;
use tokio::sync::broadcast;
use crate::ExecServerError;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecOutputStream;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::ReadParams;
use crate::protocol::ReadResponse;
use crate::protocol::TerminateResponse;
use crate::protocol::WriteResponse;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExecServerEvent {
OutputDelta(ExecOutputDeltaNotification),
Exited(ExecExitedNotification),
pub enum ExecSessionEvent {
Output {
seq: u64,
stream: ExecOutputStream,
chunk: Vec<u8>,
},
Exited {
seq: u64,
exit_code: i32,
},
Closed {
seq: u64,
},
}
#[async_trait]
pub trait ExecProcess: Send + Sync {
async fn start(&self, params: ExecParams) -> Result<ExecResponse, ExecServerError>;
fn process_id(&self) -> &str; // TODO(codex) make this a ProcessId struct
async fn read(&self, params: ReadParams) -> Result<ReadResponse, ExecServerError>;
fn subscribe(&self) -> broadcast::Receiver<ExecSessionEvent>;
async fn write(
&self,
process_id: &str,
chunk: Vec<u8>,
) -> Result<WriteResponse, ExecServerError>;
async fn write_stdin(&self, chunk: Vec<u8>) -> Result<(), ExecServerError>; // TODO(codex) rename to write()
async fn terminate(&self, process_id: &str) -> Result<TerminateResponse, ExecServerError>;
fn subscribe_events(&self) -> broadcast::Receiver<ExecServerEvent>;
async fn terminate(&self) -> Result<(), ExecServerError>;
}
#[async_trait]
pub trait ExecBackend: Send + Sync {
async fn start(&self, params: ExecParams) -> Result<Arc<dyn ExecProcess>, ExecServerError>;
}

10
codex-rs/exec-server/src/protocol.rs
View File

@@ -13,6 +13,7 @@ pub const EXEC_WRITE_METHOD: &str = "process/write";
pub const EXEC_TERMINATE_METHOD: &str = "process/terminate";
pub const EXEC_OUTPUT_DELTA_METHOD: &str = "process/output";
pub const EXEC_EXITED_METHOD: &str = "process/exited";
pub const EXEC_CLOSED_METHOD: &str = "process/closed";
pub const FS_READ_FILE_METHOD: &str = "fs/readFile";
pub const FS_WRITE_FILE_METHOD: &str = "fs/writeFile";
pub const FS_CREATE_DIRECTORY_METHOD: &str = "fs/createDirectory";
@@ -129,6 +130,7 @@ pub enum ExecOutputStream {
#[serde(rename_all = "camelCase")]
pub struct ExecOutputDeltaNotification {
pub process_id: String,
pub seq: u64,
pub stream: ExecOutputStream,
pub chunk: ByteChunk,
}
@@ -137,9 +139,17 @@ pub struct ExecOutputDeltaNotification {
#[serde(rename_all = "camelCase")]
pub struct ExecExitedNotification {
pub process_id: String,
pub seq: u64,
pub exit_code: i32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecClosedNotification {
pub process_id: String,
pub seq: u64,
}
mod base64_bytes {
use super::BASE64_STANDARD;
use base64::Engine as _;

91
codex-rs/exec-server/src/remote_process.rs
View File

@@ -1,51 +1,86 @@
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use async_trait::async_trait;
use tokio::sync::broadcast;
use crate::ExecBackend;
use crate::ExecProcess;
use crate::ExecServerClient;
use crate::ExecServerError;
use crate::ExecServerEvent;
use crate::ExecSessionEvent;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::ReadParams;
use crate::protocol::ReadResponse;
use crate::protocol::TerminateResponse;
use crate::protocol::WriteResponse;
#[derive(Clone)]
pub(crate) struct RemoteProcess {
client: ExecServerClient,
}
struct RemoteExecProcess {
process_id: String,
events: StdMutex<broadcast::Receiver<ExecSessionEvent>>,
backend: RemoteProcess,
}
impl RemoteProcess {
pub(crate) fn new(client: ExecServerClient) -> Self {
Self { client }
}
async fn write_stdin(&self, process_id: &str, chunk: Vec<u8>) -> Result<(), ExecServerError> {
let response = self.client.write(process_id, chunk).await?;
if response.accepted {
Ok(())
} else {
Err(ExecServerError::Protocol(format!(
"exec-server did not accept stdin for process {process_id}"
)))
}
}
async fn terminate_process(&self, process_id: &str) -> Result<(), ExecServerError> {
self.client.terminate(process_id).await?;
Ok(())
}
}
#[async_trait]
impl ExecProcess for RemoteProcess {
async fn start(&self, params: ExecParams) -> Result<ExecResponse, ExecServerError> {
self.client.exec(params).await
}
impl ExecBackend for RemoteProcess {
async fn start(&self, params: ExecParams) -> Result<Arc<dyn ExecProcess>, ExecServerError> {
let process_id = params.process_id.clone();
let events = self.client.register_session(&process_id).await?;
if let Err(err) = self.client.exec(params).await {
self.client.unregister_session(&process_id).await;
return Err(err);
}
async fn read(&self, params: ReadParams) -> Result<ReadResponse, ExecServerError> {
self.client.read(params).await
}
async fn write(
&self,
process_id: &str,
chunk: Vec<u8>,
) -> Result<WriteResponse, ExecServerError> {
self.client.write(process_id, chunk).await
}
async fn terminate(&self, process_id: &str) -> Result<TerminateResponse, ExecServerError> {
self.client.terminate(process_id).await
}
fn subscribe_events(&self) -> broadcast::Receiver<ExecServerEvent> {
self.client.event_receiver()
Ok(Arc::new(RemoteExecProcess {
process_id,
events: StdMutex::new(events),
backend: self.clone(),
}))
}
}
#[async_trait]
impl ExecProcess for RemoteExecProcess {
fn process_id(&self) -> &str {
&self.process_id
}
fn subscribe(&self) -> broadcast::Receiver<ExecSessionEvent> {
self
.events
.lock()
.expect("remote exec process events mutex should not be poisoned")
.resubscribe()
}
async fn write_stdin(&self, chunk: Vec<u8>) -> Result<(), ExecServerError> {
self.backend.write_stdin(&self.process_id, chunk).await
}
async fn terminate(&self) -> Result<(), ExecServerError> {
self.backend.terminate_process(&self.process_id).await
}
}

132
codex-rs/exec-server/tests/exec_process.rs
View File

@@ -8,10 +8,12 @@ use anyhow::Result;
use codex_exec_server::Environment;
use codex_exec_server::ExecParams;
use codex_exec_server::ExecProcess;
use codex_exec_server::ExecResponse;
use codex_exec_server::ReadParams;
use codex_exec_server::ExecSessionEvent;
use codex_exec_server::ExecSessionHandle;
use pretty_assertions::assert_eq;
use test_case::test_case;
use tokio::time::Duration;
use tokio::time::timeout;
use common::exec_server::ExecServerHarness;
use common::exec_server::exec_server;
@@ -40,7 +42,7 @@ async fn create_process_context(use_remote: bool) -> Result<ProcessContext> {
async fn assert_exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
let context = create_process_context(use_remote).await?;
let response = context
let mut session = context
.process
.start(ExecParams {
process_id: "proc-1".to_string(),
@@ -51,31 +53,101 @@ async fn assert_exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
arg0: None,
})
.await?;
assert_eq!(
response,
ExecResponse {
process_id: "proc-1".to_string(),
}
);
assert_eq!(session.process_id, "proc-1");
let mut next_seq = 0;
let mut exit_code = None;
loop {
let read = context
.process
.read(ReadParams {
process_id: "proc-1".to_string(),
after_seq: Some(next_seq),
max_bytes: None,
wait_ms: Some(100),
})
.await?;
next_seq = read.next_seq;
if read.exited {
assert_eq!(read.exit_code, Some(0));
break;
match timeout(Duration::from_secs(2), session.events.recv()).await?? {
ExecSessionEvent::Exited {
exit_code: code, ..
} => exit_code = Some(code),
ExecSessionEvent::Closed { .. } => break,
ExecSessionEvent::Output { .. } => {}
}
}
assert_eq!(exit_code, Some(0));
Ok(())
}
async fn collect_process_output_from_events(
mut session: ExecSessionHandle,
) -> Result<(String, i32, bool)> {
let mut output = String::new();
let mut exit_code = None;
loop {
match timeout(Duration::from_secs(2), session.events.recv()).await?? {
ExecSessionEvent::Output { chunk, .. } => {
output.push_str(&String::from_utf8_lossy(&chunk));
}
ExecSessionEvent::Exited {
exit_code: code, ..
} => exit_code = Some(code),
ExecSessionEvent::Closed { .. } => {
break;
}
}
}
Ok((output, exit_code.unwrap_or(-1), true))
}
async fn assert_exec_process_streams_output(use_remote: bool) -> Result<()> {
let context = create_process_context(use_remote).await?;
let process_id = "proc-stream".to_string();
let session = context
.process
.start(ExecParams {
process_id: process_id.clone(),
argv: vec![
"/bin/sh".to_string(),
"-c".to_string(),
"sleep 0.05; printf 'session output\\n'".to_string(),
],
cwd: std::env::current_dir()?,
env: Default::default(),
tty: false,
arg0: None,
})
.await?;
assert_eq!(session.process_id, process_id);
let (output, exit_code, closed) = collect_process_output_from_events(session).await?;
assert_eq!(output, "session output\n");
assert_eq!(exit_code, 0);
assert!(closed);
Ok(())
}
async fn assert_exec_process_write_then_read(use_remote: bool) -> Result<()> {
let context = create_process_context(use_remote).await?;
let process_id = "proc-stdin".to_string();
let session = context
.process
.start(ExecParams {
process_id: process_id.clone(),
argv: vec![
"/usr/bin/python3".to_string(),
"-c".to_string(),
"import sys; line = sys.stdin.readline(); sys.stdout.write(f'from-stdin:{line}'); sys.stdout.flush()".to_string(),
],
cwd: std::env::current_dir()?,
env: Default::default(),
tty: true,
arg0: None,
})
.await?;
assert_eq!(session.process_id, process_id);
tokio::time::sleep(Duration::from_millis(200)).await;
session.write_stdin(b"hello\n".to_vec()).await?;
let (output, exit_code, closed) = collect_process_output_from_events(session).await?;
assert!(
output.contains("from-stdin:hello"),
"unexpected output: {output:?}"
);
assert_eq!(exit_code, 0);
assert!(closed);
Ok(())
}
@@ -85,3 +157,17 @@ async fn assert_exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
async fn exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
assert_exec_process_starts_and_exits(use_remote).await
}
#[test_case(false ; "local")]
#[test_case(true ; "remote")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_process_streams_output(use_remote: bool) -> Result<()> {
assert_exec_process_streams_output(use_remote).await
}
#[test_case(false ; "local")]
#[test_case(true ; "remote")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_process_write_then_read(use_remote: bool) -> Result<()> {
assert_exec_process_write_then_read(use_remote).await
}