Compare commits

...

1 Commits

Author SHA1 Message Date
starr-openai
fd58e4621f Add exec-server crate and tests
Co-authored-by: Codex <noreply@openai.com>
2026-03-16 17:30:51 -07:00
12 changed files with 1314 additions and 1 deletions

17
codex-rs/Cargo.lock generated
View File

@@ -1989,6 +1989,23 @@ dependencies = [
"wiremock",
]
[[package]]
name = "codex-exec-server"
version = "0.0.0"
dependencies = [
"anyhow",
"base64 0.22.1",
"codex-app-server-protocol",
"codex-utils-cargo-bin",
"codex-utils-pty",
"pretty_assertions",
"serde",
"serde_json",
"thiserror 2.0.18",
"tokio",
"tracing",
]
[[package]]
name = "codex-execpolicy"
version = "0.0.0"

View File

@@ -25,6 +25,7 @@ members = [
"hooks",
"secrets",
"exec",
"exec-server",
"execpolicy",
"execpolicy-legacy",
"keyring-store",

View File

@@ -0,0 +1,35 @@
[package]
name = "codex-exec-server"
version.workspace = true
edition.workspace = true
license.workspace = true
[[bin]]
name = "codex-exec-server"
path = "src/bin/codex-exec-server.rs"
[lints]
workspace = true
[dependencies]
base64 = { workspace = true }
codex-app-server-protocol = { workspace = true }
codex-utils-pty = { workspace = true }
serde = { workspace = true, features = ["derive"] }
serde_json = { workspace = true }
thiserror = { workspace = true }
tokio = { workspace = true, features = [
"io-std",
"io-util",
"macros",
"process",
"rt-multi-thread",
"sync",
"time",
] }
tracing = { workspace = true }
[dev-dependencies]
anyhow = { workspace = true }
codex-utils-cargo-bin = { workspace = true }
pretty_assertions = { workspace = true }

View File

@@ -0,0 +1,7 @@
#[tokio::main]
async fn main() {
if let Err(err) = codex_exec_server::run_main().await {
eprintln!("{err}");
std::process::exit(1);
}
}

View File

@@ -0,0 +1,514 @@
use std::collections::HashMap;
use std::path::PathBuf;
use std::process::Stdio;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::AtomicI64;
use std::sync::atomic::Ordering;
use codex_app_server_protocol::JSONRPCError;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCRequest;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use serde::Serialize;
use serde::de::DeserializeOwned;
use serde_json::Value;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::process::Child;
use tokio::process::Command;
use tokio::sync::Mutex;
use tokio::sync::broadcast;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
use tracing::debug;
use tracing::warn;
use crate::protocol::EXEC_EXITED_METHOD;
use crate::protocol::EXEC_METHOD;
use crate::protocol::EXEC_OUTPUT_DELTA_METHOD;
use crate::protocol::EXEC_TERMINATE_METHOD;
use crate::protocol::EXEC_WRITE_METHOD;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeParams;
use crate::protocol::InitializeResponse;
use crate::protocol::TerminateParams;
use crate::protocol::TerminateResponse;
use crate::protocol::WriteParams;
use crate::protocol::WriteResponse;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExecServerLaunchCommand {
pub program: PathBuf,
pub args: Vec<String>,
}
pub struct ExecServerProcess {
process_id: String,
pid: Option<u32>,
output_rx: Option<broadcast::Receiver<Vec<u8>>>,
writer_tx: mpsc::Sender<Vec<u8>>,
status: Arc<RemoteProcessStatus>,
client: ExecServerClient,
}
impl std::fmt::Debug for ExecServerProcess {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("ExecServerProcess")
.field("process_id", &self.process_id)
.field("pid", &self.pid)
.field("has_exited", &self.has_exited())
.field("exit_code", &self.exit_code())
.finish()
}
}
impl ExecServerProcess {
pub fn writer_sender(&self) -> mpsc::Sender<Vec<u8>> {
self.writer_tx.clone()
}
pub fn output_receiver(&self) -> broadcast::Receiver<Vec<u8>> {
match self.output_rx.as_ref() {
Some(output_rx) => output_rx.resubscribe(),
None => panic!("output receiver should still be present"),
}
}
pub fn take_output_receiver(&mut self) -> broadcast::Receiver<Vec<u8>> {
match self.output_rx.take() {
Some(output_rx) => output_rx,
None => panic!("output receiver should only be taken once"),
}
}
pub fn has_exited(&self) -> bool {
self.status.has_exited()
}
pub fn exit_code(&self) -> Option<i32> {
self.status.exit_code()
}
pub fn pid(&self) -> Option<u32> {
self.pid
}
pub fn terminate(&self) {
self.status.mark_exited(None);
let client = self.client.clone();
let process_id = self.process_id.clone();
tokio::spawn(async move {
let _ = client.terminate_process(&process_id).await;
});
}
}
impl std::fmt::Debug for RemoteProcessStatus {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("RemoteProcessStatus")
.field("exited", &self.has_exited())
.field("exit_code", &self.exit_code())
.finish()
}
}
struct RemoteProcessStatus {
exited: AtomicBool,
exit_code: StdMutex<Option<i32>>,
}
impl RemoteProcessStatus {
fn new() -> Self {
Self {
exited: AtomicBool::new(false),
exit_code: StdMutex::new(None),
}
}
fn has_exited(&self) -> bool {
self.exited.load(Ordering::SeqCst)
}
fn exit_code(&self) -> Option<i32> {
self.exit_code.lock().ok().and_then(|guard| *guard)
}
fn mark_exited(&self, exit_code: Option<i32>) {
self.exited.store(true, Ordering::SeqCst);
if let Ok(mut guard) = self.exit_code.lock() {
*guard = exit_code;
}
}
}
struct RegisteredProcess {
output_tx: broadcast::Sender<Vec<u8>>,
status: Arc<RemoteProcessStatus>,
}
struct Inner {
child: StdMutex<Option<Child>>,
write_tx: mpsc::UnboundedSender<JSONRPCMessage>,
pending: Mutex<HashMap<RequestId, oneshot::Sender<Result<Value, JSONRPCErrorError>>>>,
processes: Mutex<HashMap<String, RegisteredProcess>>,
next_request_id: AtomicI64,
reader_task: JoinHandle<()>,
writer_task: JoinHandle<()>,
}
impl Drop for Inner {
fn drop(&mut self) {
self.reader_task.abort();
self.writer_task.abort();
if let Ok(mut child_guard) = self.child.lock()
&& let Some(child) = child_guard.as_mut()
{
let _ = child.start_kill();
}
}
}
#[derive(Clone)]
pub struct ExecServerClient {
inner: Arc<Inner>,
}
#[derive(Debug, thiserror::Error)]
pub enum ExecServerError {
#[error("failed to spawn exec-server: {0}")]
Spawn(#[source] std::io::Error),
#[error("exec-server transport closed")]
Closed,
#[error("failed to serialize or deserialize exec-server JSON: {0}")]
Json(#[from] serde_json::Error),
#[error("exec-server protocol error: {0}")]
Protocol(String),
#[error("exec-server rejected request ({code}): {message}")]
Server { code: i64, message: String },
}
impl ExecServerClient {
pub async fn spawn(command: ExecServerLaunchCommand) -> Result<Self, ExecServerError> {
let mut child = Command::new(&command.program);
child.args(&command.args);
child.stdin(Stdio::piped());
child.stdout(Stdio::piped());
child.stderr(Stdio::inherit());
child.kill_on_drop(true);
let mut child = child.spawn().map_err(ExecServerError::Spawn)?;
let stdin = child.stdin.take().ok_or_else(|| {
ExecServerError::Protocol("exec-server stdin was not captured".to_string())
})?;
let stdout = child.stdout.take().ok_or_else(|| {
ExecServerError::Protocol("exec-server stdout was not captured".to_string())
})?;
let (write_tx, mut write_rx) = mpsc::unbounded_channel::<JSONRPCMessage>();
let writer_task = tokio::spawn(async move {
let mut stdin = stdin;
while let Some(message) = write_rx.recv().await {
let encoded = match serde_json::to_vec(&message) {
Ok(encoded) => encoded,
Err(err) => {
warn!("failed to encode exec-server message: {err}");
break;
}
};
if stdin.write_all(&encoded).await.is_err() {
break;
}
if stdin.write_all(b"\n").await.is_err() {
break;
}
if stdin.flush().await.is_err() {
break;
}
}
});
let pending = Mutex::new(HashMap::<
RequestId,
oneshot::Sender<Result<Value, JSONRPCErrorError>>,
>::new());
let processes = Mutex::new(HashMap::<String, RegisteredProcess>::new());
let inner = Arc::new_cyclic(move |weak| {
let weak = weak.clone();
let reader_task = tokio::spawn(async move {
let mut lines = BufReader::new(stdout).lines();
loop {
let Some(inner) = weak.upgrade() else {
break;
};
let next_line = lines.next_line().await;
match next_line {
Ok(Some(line)) => {
if line.trim().is_empty() {
continue;
}
match serde_json::from_str::<JSONRPCMessage>(&line) {
Ok(message) => {
if let Err(err) = handle_server_message(&inner, message).await {
warn!("failed to handle exec-server message: {err}");
break;
}
}
Err(err) => {
warn!("failed to parse exec-server message: {err}");
break;
}
}
}
Ok(None) => break,
Err(err) => {
warn!("failed to read exec-server stdout: {err}");
break;
}
}
}
if let Some(inner) = weak.upgrade() {
handle_transport_shutdown(&inner).await;
}
});
Inner {
child: StdMutex::new(Some(child)),
write_tx,
pending,
processes,
next_request_id: AtomicI64::new(1),
reader_task,
writer_task,
}
});
let client = Self { inner };
client.initialize().await?;
Ok(client)
}
pub async fn start_process(
&self,
params: ExecParams,
) -> Result<ExecServerProcess, ExecServerError> {
let process_id = params.process_id.clone();
let status = Arc::new(RemoteProcessStatus::new());
let (output_tx, output_rx) = broadcast::channel(256);
self.inner.processes.lock().await.insert(
process_id.clone(),
RegisteredProcess {
output_tx,
status: Arc::clone(&status),
},
);
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
let client = self.clone();
let write_process_id = process_id.clone();
tokio::spawn(async move {
while let Some(chunk) = writer_rx.recv().await {
let request = WriteParams {
process_id: write_process_id.clone(),
chunk: chunk.into(),
};
if client.write_process(request).await.is_err() {
break;
}
}
});
let response = match self.request::<_, ExecResponse>(EXEC_METHOD, &params).await {
Ok(response) => response,
Err(err) => {
self.inner.processes.lock().await.remove(&process_id);
return Err(err);
}
};
if let Some(exit_code) = response.exit_code {
status.mark_exited(Some(exit_code));
}
Ok(ExecServerProcess {
process_id,
pid: response.pid,
output_rx: Some(output_rx),
writer_tx,
status,
client: self.clone(),
})
}
async fn initialize(&self) -> Result<(), ExecServerError> {
let _: InitializeResponse = self
.request(
INITIALIZE_METHOD,
&InitializeParams {
client_name: "codex-core".to_string(),
},
)
.await?;
self.notify(INITIALIZED_METHOD, &serde_json::json!({}))
.await
}
async fn write_process(&self, params: WriteParams) -> Result<WriteResponse, ExecServerError> {
self.request(EXEC_WRITE_METHOD, &params).await
}
async fn terminate_process(
&self,
process_id: &str,
) -> Result<TerminateResponse, ExecServerError> {
self.request(
EXEC_TERMINATE_METHOD,
&TerminateParams {
process_id: process_id.to_string(),
},
)
.await
}
async fn notify<P: Serialize>(&self, method: &str, params: &P) -> Result<(), ExecServerError> {
let params = serde_json::to_value(params)?;
self.inner
.write_tx
.send(JSONRPCMessage::Notification(JSONRPCNotification {
method: method.to_string(),
params: Some(params),
}))
.map_err(|_| ExecServerError::Closed)
}
async fn request<P, R>(&self, method: &str, params: &P) -> Result<R, ExecServerError>
where
P: Serialize,
R: DeserializeOwned,
{
let request_id =
RequestId::Integer(self.inner.next_request_id.fetch_add(1, Ordering::SeqCst));
let (response_tx, response_rx) = oneshot::channel();
self.inner
.pending
.lock()
.await
.insert(request_id.clone(), response_tx);
let params = serde_json::to_value(params)?;
let message = JSONRPCMessage::Request(JSONRPCRequest {
id: request_id.clone(),
method: method.to_string(),
params: Some(params),
trace: None,
});
if self.inner.write_tx.send(message).is_err() {
self.inner.pending.lock().await.remove(&request_id);
return Err(ExecServerError::Closed);
}
let result = response_rx.await.map_err(|_| ExecServerError::Closed)?;
match result {
Ok(value) => serde_json::from_value(value).map_err(ExecServerError::from),
Err(error) => Err(ExecServerError::Server {
code: error.code,
message: error.message,
}),
}
}
}
async fn handle_server_message(
inner: &Arc<Inner>,
message: JSONRPCMessage,
) -> Result<(), ExecServerError> {
match message {
JSONRPCMessage::Response(JSONRPCResponse { id, result }) => {
if let Some(tx) = inner.pending.lock().await.remove(&id) {
let _ = tx.send(Ok(result));
}
}
JSONRPCMessage::Error(JSONRPCError { id, error }) => {
if let Some(tx) = inner.pending.lock().await.remove(&id) {
let _ = tx.send(Err(error));
}
}
JSONRPCMessage::Notification(notification) => {
handle_server_notification(inner, notification).await?;
}
JSONRPCMessage::Request(request) => {
return Err(ExecServerError::Protocol(format!(
"unexpected exec-server request from child: {}",
request.method
)));
}
}
Ok(())
}
async fn handle_server_notification(
inner: &Arc<Inner>,
notification: JSONRPCNotification,
) -> Result<(), ExecServerError> {
match notification.method.as_str() {
EXEC_OUTPUT_DELTA_METHOD => {
let params: ExecOutputDeltaNotification =
serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
let chunk = params.chunk.into_inner();
let processes = inner.processes.lock().await;
if let Some(process) = processes.get(&params.process_id) {
let _ = process.output_tx.send(chunk);
}
}
EXEC_EXITED_METHOD => {
let params: ExecExitedNotification =
serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
let mut processes = inner.processes.lock().await;
if let Some(process) = processes.remove(&params.process_id) {
process.status.mark_exited(Some(params.exit_code));
}
}
other => {
debug!("ignoring unknown exec-server notification: {other}");
}
}
Ok(())
}
async fn handle_transport_shutdown(inner: &Arc<Inner>) {
let pending = {
let mut pending = inner.pending.lock().await;
pending.drain().map(|(_, tx)| tx).collect::<Vec<_>>()
};
for tx in pending {
let _ = tx.send(Err(JSONRPCErrorError {
code: -32000,
data: None,
message: "exec-server transport closed".to_string(),
}));
}
let processes = {
let mut processes = inner.processes.lock().await;
processes
.drain()
.map(|(_, process)| process)
.collect::<Vec<_>>()
};
for process in processes {
process.status.mark_exited(None);
}
}

View File

@@ -0,0 +1,20 @@
mod client;
mod protocol;
mod server;
pub use client::ExecServerClient;
pub use client::ExecServerError;
pub use client::ExecServerLaunchCommand;
pub use client::ExecServerProcess;
pub use protocol::ExecExitedNotification;
pub use protocol::ExecOutputDeltaNotification;
pub use protocol::ExecOutputStream;
pub use protocol::ExecParams;
pub use protocol::ExecResponse;
pub use protocol::InitializeParams;
pub use protocol::InitializeResponse;
pub use protocol::TerminateParams;
pub use protocol::TerminateResponse;
pub use protocol::WriteParams;
pub use protocol::WriteResponse;
pub use server::run_main;

View File

@@ -0,0 +1,144 @@
use std::collections::HashMap;
use std::path::PathBuf;
use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use codex_utils_pty::DEFAULT_OUTPUT_BYTES_CAP;
use serde::Deserialize;
use serde::Serialize;
pub const INITIALIZE_METHOD: &str = "initialize";
pub const INITIALIZED_METHOD: &str = "initialized";
pub const EXEC_METHOD: &str = "command/exec";
pub const EXEC_WRITE_METHOD: &str = "command/exec/write";
pub const EXEC_TERMINATE_METHOD: &str = "command/exec/terminate";
pub const EXEC_OUTPUT_DELTA_METHOD: &str = "command/exec/outputDelta";
pub const EXEC_EXITED_METHOD: &str = "command/exec/exited";
pub const PROTOCOL_VERSION: &str = "exec-server.v0";
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ByteChunk(#[serde(with = "base64_bytes")] pub Vec<u8>);
impl ByteChunk {
pub fn into_inner(self) -> Vec<u8> {
self.0
}
}
impl From<Vec<u8>> for ByteChunk {
fn from(value: Vec<u8>) -> Self {
Self(value)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InitializeParams {
pub client_name: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct InitializeResponse {
pub protocol_version: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecParams {
pub process_id: String,
pub argv: Vec<String>,
pub cwd: PathBuf,
pub env: HashMap<String, String>,
pub tty: bool,
#[serde(default = "default_output_bytes_cap")]
pub output_bytes_cap: usize,
pub arg0: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecResponse {
pub process_id: String,
pub pid: Option<u32>,
pub running: bool,
pub exit_code: Option<i32>,
pub stdout: Option<ByteChunk>,
pub stderr: Option<ByteChunk>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WriteParams {
pub process_id: String,
pub chunk: ByteChunk,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WriteResponse {
pub accepted: bool,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TerminateParams {
pub process_id: String,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TerminateResponse {
pub running: bool,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum ExecOutputStream {
Stdout,
Stderr,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecOutputDeltaNotification {
pub process_id: String,
pub stream: ExecOutputStream,
pub chunk: ByteChunk,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecExitedNotification {
pub process_id: String,
pub exit_code: i32,
}
fn default_output_bytes_cap() -> usize {
DEFAULT_OUTPUT_BYTES_CAP
}
mod base64_bytes {
use super::BASE64_STANDARD;
use base64::Engine as _;
use serde::Deserialize;
use serde::Deserializer;
use serde::Serializer;
pub fn serialize<S>(bytes: &[u8], serializer: S) -> Result<S::Ok, S::Error>
where
S: Serializer,
{
serializer.serialize_str(&BASE64_STANDARD.encode(bytes))
}
pub fn deserialize<'de, D>(deserializer: D) -> Result<Vec<u8>, D::Error>
where
D: Deserializer<'de>,
{
let encoded = String::deserialize(deserializer)?;
BASE64_STANDARD
.decode(encoded)
.map_err(serde::de::Error::custom)
}
}

View File

@@ -0,0 +1,422 @@
use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::Arc;
use std::sync::Mutex as StdMutex;
use codex_app_server_protocol::JSONRPCError;
use codex_app_server_protocol::JSONRPCErrorError;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCRequest;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_utils_pty::ExecCommandSession;
use codex_utils_pty::TerminalSize;
use serde::Serialize;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::io::BufWriter;
use tokio::sync::Mutex;
use crate::protocol::EXEC_EXITED_METHOD;
use crate::protocol::EXEC_METHOD;
use crate::protocol::EXEC_OUTPUT_DELTA_METHOD;
use crate::protocol::EXEC_TERMINATE_METHOD;
use crate::protocol::EXEC_WRITE_METHOD;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecOutputStream;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeResponse;
use crate::protocol::PROTOCOL_VERSION;
use crate::protocol::TerminateParams;
use crate::protocol::TerminateResponse;
use crate::protocol::WriteParams;
use crate::protocol::WriteResponse;
struct RunningProcess {
session: ExecCommandSession,
tty: bool,
stdout_buffer: Arc<StdMutex<BoundedBytesBuffer>>,
stderr_buffer: Arc<StdMutex<BoundedBytesBuffer>>,
}
#[derive(Debug)]
struct BoundedBytesBuffer {
max_bytes: usize,
bytes: VecDeque<u8>,
}
impl BoundedBytesBuffer {
fn new(max_bytes: usize) -> Self {
Self {
max_bytes,
bytes: VecDeque::with_capacity(max_bytes.min(8192)),
}
}
fn push_chunk(&mut self, chunk: &[u8]) {
if self.max_bytes == 0 {
return;
}
for byte in chunk {
self.bytes.push_back(*byte);
if self.bytes.len() > self.max_bytes {
self.bytes.pop_front();
}
}
}
fn snapshot(&self) -> Vec<u8> {
self.bytes.iter().copied().collect()
}
}
pub async fn run_main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let writer = Arc::new(Mutex::new(BufWriter::new(tokio::io::stdout())));
let processes = Arc::new(Mutex::new(HashMap::<String, RunningProcess>::new()));
let mut lines = BufReader::new(tokio::io::stdin()).lines();
while let Some(line) = lines.next_line().await? {
if line.trim().is_empty() {
continue;
}
let message = serde_json::from_str::<JSONRPCMessage>(&line)?;
if let JSONRPCMessage::Request(request) = message {
handle_request(request, &writer, &processes).await;
continue;
}
if let JSONRPCMessage::Notification(notification) = message {
if notification.method != INITIALIZED_METHOD {
send_error(
&writer,
RequestId::Integer(-1),
invalid_request(format!(
"unexpected notification method: {}",
notification.method
)),
)
.await;
}
continue;
}
}
let remaining = {
let mut processes = processes.lock().await;
processes
.drain()
.map(|(_, process)| process)
.collect::<Vec<_>>()
};
for process in remaining {
process.session.terminate();
}
Ok(())
}
async fn handle_request(
request: JSONRPCRequest,
writer: &Arc<Mutex<BufWriter<tokio::io::Stdout>>>,
processes: &Arc<Mutex<HashMap<String, RunningProcess>>>,
) {
let response = match request.method.as_str() {
INITIALIZE_METHOD => serde_json::to_value(InitializeResponse {
protocol_version: PROTOCOL_VERSION.to_string(),
})
.map_err(|err| internal_error(err.to_string())),
EXEC_METHOD => handle_exec_request(request.params, writer, processes).await,
EXEC_WRITE_METHOD => handle_write_request(request.params, processes).await,
EXEC_TERMINATE_METHOD => handle_terminate_request(request.params, processes).await,
other => Err(invalid_request(format!("unknown method: {other}"))),
};
match response {
Ok(result) => {
send_response(
writer,
JSONRPCResponse {
id: request.id,
result,
},
)
.await;
}
Err(err) => {
send_error(writer, request.id, err).await;
}
}
}
async fn handle_exec_request(
params: Option<serde_json::Value>,
writer: &Arc<Mutex<BufWriter<tokio::io::Stdout>>>,
processes: &Arc<Mutex<HashMap<String, RunningProcess>>>,
) -> Result<serde_json::Value, JSONRPCErrorError> {
let params: ExecParams = serde_json::from_value(params.unwrap_or(serde_json::Value::Null))
.map_err(|err| invalid_params(err.to_string()))?;
let (program, args) = params
.argv
.split_first()
.ok_or_else(|| invalid_params("argv must not be empty".to_string()))?;
let spawned = if params.tty {
codex_utils_pty::spawn_pty_process(
program,
args,
params.cwd.as_path(),
&params.env,
&params.arg0,
TerminalSize::default(),
)
.await
} else {
codex_utils_pty::spawn_pipe_process_no_stdin(
program,
args,
params.cwd.as_path(),
&params.env,
&params.arg0,
)
.await
}
.map_err(|err| internal_error(err.to_string()))?;
let pid = spawned.session.pid();
let stdout_buffer = Arc::new(StdMutex::new(BoundedBytesBuffer::new(
params.output_bytes_cap,
)));
let stderr_buffer = Arc::new(StdMutex::new(BoundedBytesBuffer::new(
params.output_bytes_cap,
)));
let process_id = params.process_id.clone();
{
let mut process_map = processes.lock().await;
if process_map.contains_key(&process_id) {
spawned.session.terminate();
return Err(invalid_request(format!(
"process {} already exists",
params.process_id
)));
}
process_map.insert(
process_id.clone(),
RunningProcess {
session: spawned.session,
tty: params.tty,
stdout_buffer: Arc::clone(&stdout_buffer),
stderr_buffer: Arc::clone(&stderr_buffer),
},
);
}
tokio::spawn(stream_output(
process_id.clone(),
ExecOutputStream::Stdout,
spawned.stdout_rx,
Arc::clone(writer),
Arc::clone(&stdout_buffer),
));
tokio::spawn(stream_output(
process_id.clone(),
ExecOutputStream::Stderr,
spawned.stderr_rx,
Arc::clone(writer),
Arc::clone(&stderr_buffer),
));
tokio::spawn(watch_exit(
process_id.clone(),
spawned.exit_rx,
Arc::clone(writer),
Arc::clone(processes),
));
serde_json::to_value(ExecResponse {
process_id,
pid,
running: true,
exit_code: None,
stdout: None,
stderr: None,
})
.map_err(|err| internal_error(err.to_string()))
}
async fn handle_write_request(
params: Option<serde_json::Value>,
processes: &Arc<Mutex<HashMap<String, RunningProcess>>>,
) -> Result<serde_json::Value, JSONRPCErrorError> {
let params: WriteParams = serde_json::from_value(params.unwrap_or(serde_json::Value::Null))
.map_err(|err| invalid_params(err.to_string()))?;
let writer_tx = {
let process_map = processes.lock().await;
let process = process_map
.get(&params.process_id)
.ok_or_else(|| invalid_request(format!("unknown process id {}", params.process_id)))?;
if !process.tty {
return Err(invalid_request(format!(
"stdin is closed for process {}",
params.process_id
)));
}
process.session.writer_sender()
};
writer_tx
.send(params.chunk.into_inner())
.await
.map_err(|_| internal_error("failed to write to process stdin".to_string()))?;
serde_json::to_value(WriteResponse { accepted: true })
.map_err(|err| internal_error(err.to_string()))
}
async fn handle_terminate_request(
params: Option<serde_json::Value>,
processes: &Arc<Mutex<HashMap<String, RunningProcess>>>,
) -> Result<serde_json::Value, JSONRPCErrorError> {
let params: TerminateParams = serde_json::from_value(params.unwrap_or(serde_json::Value::Null))
.map_err(|err| invalid_params(err.to_string()))?;
let process = {
let mut process_map = processes.lock().await;
process_map.remove(&params.process_id)
};
if let Some(process) = process {
process.session.terminate();
serde_json::to_value(TerminateResponse { running: true })
.map_err(|err| internal_error(err.to_string()))
} else {
serde_json::to_value(TerminateResponse { running: false })
.map_err(|err| internal_error(err.to_string()))
}
}
async fn stream_output(
process_id: String,
stream: ExecOutputStream,
mut receiver: tokio::sync::mpsc::Receiver<Vec<u8>>,
writer: Arc<Mutex<BufWriter<tokio::io::Stdout>>>,
buffer: Arc<StdMutex<BoundedBytesBuffer>>,
) {
while let Some(chunk) = receiver.recv().await {
if let Ok(mut guard) = buffer.lock() {
guard.push_chunk(&chunk);
}
let notification = ExecOutputDeltaNotification {
process_id: process_id.clone(),
stream,
chunk: chunk.into(),
};
if send_notification(&writer, EXEC_OUTPUT_DELTA_METHOD, &notification)
.await
.is_err()
{
break;
}
}
}
async fn watch_exit(
process_id: String,
exit_rx: tokio::sync::oneshot::Receiver<i32>,
writer: Arc<Mutex<BufWriter<tokio::io::Stdout>>>,
processes: Arc<Mutex<HashMap<String, RunningProcess>>>,
) {
let exit_code = exit_rx.await.unwrap_or(-1);
let removed = {
let mut processes = processes.lock().await;
processes.remove(&process_id)
};
if let Some(process) = removed {
let _ = process.stdout_buffer.lock().map(|buffer| buffer.snapshot());
let _ = process.stderr_buffer.lock().map(|buffer| buffer.snapshot());
}
let _ = send_notification(
&writer,
EXEC_EXITED_METHOD,
&ExecExitedNotification {
process_id,
exit_code,
},
)
.await;
}
async fn send_response(
writer: &Arc<Mutex<BufWriter<tokio::io::Stdout>>>,
response: JSONRPCResponse,
) {
let _ = send_message(writer, JSONRPCMessage::Response(response)).await;
}
async fn send_error(
writer: &Arc<Mutex<BufWriter<tokio::io::Stdout>>>,
id: RequestId,
error: JSONRPCErrorError,
) {
let _ = send_message(writer, JSONRPCMessage::Error(JSONRPCError { error, id })).await;
}
async fn send_notification<T: Serialize>(
writer: &Arc<Mutex<BufWriter<tokio::io::Stdout>>>,
method: &str,
params: &T,
) -> Result<(), serde_json::Error> {
send_message(
writer,
JSONRPCMessage::Notification(JSONRPCNotification {
method: method.to_string(),
params: Some(serde_json::to_value(params)?),
}),
)
.await
.map_err(serde_json::Error::io)
}
async fn send_message(
writer: &Arc<Mutex<BufWriter<tokio::io::Stdout>>>,
message: JSONRPCMessage,
) -> std::io::Result<()> {
let encoded =
serde_json::to_vec(&message).map_err(|err| std::io::Error::other(err.to_string()))?;
let mut writer = writer.lock().await;
writer.write_all(&encoded).await?;
writer.write_all(b"\n").await?;
writer.flush().await
}
fn invalid_request(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32600,
data: None,
message,
}
}
fn invalid_params(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32602,
data: None,
message,
}
}
fn internal_error(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32603,
data: None,
message,
}
}

View File

@@ -0,0 +1,141 @@
#![cfg(unix)]
use std::process::Stdio;
use std::time::Duration;
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCRequest;
use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use codex_exec_server::ExecParams;
use codex_exec_server::ExecServerClient;
use codex_exec_server::ExecServerLaunchCommand;
use codex_exec_server::InitializeParams;
use codex_exec_server::InitializeResponse;
use codex_utils_cargo_bin::cargo_bin;
use pretty_assertions::assert_eq;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::process::Command;
use tokio::sync::broadcast;
use tokio::time::timeout;
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_accepts_initialize_over_stdio() -> anyhow::Result<()> {
let binary = cargo_bin("codex-exec-server")?;
let mut child = Command::new(binary);
child.stdin(Stdio::piped());
child.stdout(Stdio::piped());
child.stderr(Stdio::inherit());
let mut child = child.spawn()?;
let mut stdin = child.stdin.take().expect("stdin");
let stdout = child.stdout.take().expect("stdout");
let mut stdout = BufReader::new(stdout).lines();
let initialize = JSONRPCMessage::Request(JSONRPCRequest {
id: RequestId::Integer(1),
method: "initialize".to_string(),
params: Some(serde_json::to_value(InitializeParams {
client_name: "exec-server-test".to_string(),
})?),
trace: None,
});
stdin
.write_all(format!("{}\n", serde_json::to_string(&initialize)?).as_bytes())
.await?;
let response_line = timeout(Duration::from_secs(5), stdout.next_line()).await??;
let response_line = response_line.expect("response line");
let response: JSONRPCMessage = serde_json::from_str(&response_line)?;
let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = response else {
panic!("expected initialize response");
};
assert_eq!(id, RequestId::Integer(1));
let initialize_response: InitializeResponse = serde_json::from_value(result)?;
assert_eq!(initialize_response.protocol_version, "exec-server.v0");
let initialized = JSONRPCMessage::Notification(JSONRPCNotification {
method: "initialized".to_string(),
params: Some(serde_json::json!({})),
});
stdin
.write_all(format!("{}\n", serde_json::to_string(&initialized)?).as_bytes())
.await?;
child.start_kill()?;
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_client_streams_output_and_accepts_writes() -> anyhow::Result<()> {
let mut env = std::collections::HashMap::new();
if let Some(path) = std::env::var_os("PATH") {
env.insert("PATH".to_string(), path.to_string_lossy().into_owned());
}
let client = ExecServerClient::spawn(ExecServerLaunchCommand {
program: cargo_bin("codex-exec-server")?,
args: Vec::new(),
})
.await?;
let process = client
.start_process(ExecParams {
process_id: "2001".to_string(),
argv: vec![
"bash".to_string(),
"-lc".to_string(),
"printf 'ready\\n'; while IFS= read -r line; do printf 'echo:%s\\n' \"$line\"; done"
.to_string(),
],
cwd: std::env::current_dir()?,
env,
tty: true,
output_bytes_cap: 4096,
arg0: None,
})
.await?;
let mut output = process.output_receiver();
assert!(
recv_until_contains(&mut output, "ready")
.await?
.contains("ready"),
"expected initial ready output"
);
process
.writer_sender()
.send(b"hello\n".to_vec())
.await
.expect("write should succeed");
assert!(
recv_until_contains(&mut output, "echo:hello")
.await?
.contains("echo:hello"),
"expected echoed output"
);
process.terminate();
Ok(())
}
async fn recv_until_contains(
output: &mut broadcast::Receiver<Vec<u8>>,
needle: &str,
) -> anyhow::Result<String> {
let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
let mut collected = String::new();
loop {
let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
let chunk = timeout(remaining, output.recv()).await??;
collected.push_str(&String::from_utf8_lossy(&chunk));
if collected.contains(needle) {
return Ok(collected);
}
}
}

View File

@@ -238,6 +238,7 @@ async fn spawn_process_with_stdin_mode(
wait_handle,
exit_status,
exit_code,
Some(pid),
None,
);

View File

@@ -79,6 +79,7 @@ pub struct ProcessHandle {
wait_handle: StdMutex<Option<JoinHandle<()>>>,
exit_status: Arc<AtomicBool>,
exit_code: Arc<StdMutex<Option<i32>>>,
pid: Option<u32>,
// PtyHandles must be preserved because the process will receive Control+C if the
// slave is closed
_pty_handles: StdMutex<Option<PtyHandles>>,
@@ -101,6 +102,7 @@ impl ProcessHandle {
wait_handle: JoinHandle<()>,
exit_status: Arc<AtomicBool>,
exit_code: Arc<StdMutex<Option<i32>>>,
pid: Option<u32>,
pty_handles: Option<PtyHandles>,
) -> Self {
Self {
@@ -112,6 +114,7 @@ impl ProcessHandle {
wait_handle: StdMutex::new(Some(wait_handle)),
exit_status,
exit_code,
pid,
_pty_handles: StdMutex::new(pty_handles),
}
}
@@ -139,6 +142,11 @@ impl ProcessHandle {
self.exit_code.lock().ok().and_then(|guard| *guard)
}
/// Returns the OS process ID when known.
pub fn pid(&self) -> Option<u32> {
self.pid
}
/// Resize the PTY in character cells.
pub fn resize(&self, size: TerminalSize) -> anyhow::Result<()> {
let handles = self

View File

@@ -159,11 +159,12 @@ async fn spawn_process_portable(
}
let mut child = pair.slave.spawn_command(command_builder)?;
let pid = child.process_id();
#[cfg(unix)]
// portable-pty establishes the spawned PTY child as a new session leader on
// Unix, so PID == PGID and we can reuse the pipe backend's process-group
// hard-kill semantics for descendants.
let process_group_id = child.process_id();
let process_group_id = pid;
let killer = child.clone_killer();
let (writer_tx, mut writer_rx) = mpsc::channel::<Vec<u8>>(128);
@@ -241,6 +242,7 @@ async fn spawn_process_portable(
wait_handle,
exit_status,
exit_code,
pid,
Some(handles),
);
@@ -394,6 +396,7 @@ async fn spawn_process_preserving_fds(
wait_handle,
exit_status,
exit_code,
Some(process_group_id),
Some(handles),
);