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

Keep first exec-server PR initialize-only

Browse Source 
Move process-oriented client helpers and exports into the exec follow-up. Keep the first PR focused on spawning the stub server and running the initialize handshake.\n\nCo-authored-by: Codex <noreply@openai.com>
This commit is contained in:
starr-openai
2026-03-18 13:01:16 -07:00
parent 2958067cf9
commit 76071974bb
4 changed files with 130 additions and 439 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -1,9 +1,6 @@
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;
@@ -21,132 +18,23 @@ 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,
output_rx: 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("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>> {
self.output_rx.resubscribe()
}
pub fn has_exited(&self) -> bool {
self.status.has_exited()
}
pub fn exit_code(&self) -> Option<i32> {
self.status.exit_code()
}
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>,
}
use crate::server_process::ExecServerLaunchCommand;
use crate::server_process::spawn_stdio_exec_server;
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<()>,
@@ -185,20 +73,11 @@ pub enum ExecServerError {
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 crate::server_process::SpawnedStdioExecServer {
child,
stdin,
stdout,
} = spawn_stdio_exec_server(command)?;
let (write_tx, mut write_rx) = mpsc::unbounded_channel::<JSONRPCMessage>();
let writer_task = tokio::spawn(async move {
@@ -227,7 +106,6 @@ impl ExecServerClient {
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 {
@@ -236,12 +114,12 @@ impl ExecServerClient {
let Some(inner) = weak.upgrade() else {
break;
};
let next_line = lines.next_line().await;
match next_line {
match lines.next_line().await {
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 {
@@ -262,8 +140,9 @@ impl ExecServerClient {
}
}
}
if let Some(inner) = weak.upgrade() {
handle_transport_shutdown(&inner).await;
fail_pending_requests(&inner).await;
}
});
@@ -271,7 +150,6 @@ impl ExecServerClient {
child: StdMutex::new(Some(child)),
write_tx,
pending,
processes,
next_request_id: AtomicI64::new(1),
reader_task,
writer_task,
@@ -279,92 +157,56 @@ impl ExecServerClient {
});
let client = Self { inner };
client.initialize().await?;
client
.initialize(InitializeParams {
client_name: "codex-core".to_string(),
})
.await?;
client.send_notification(INITIALIZED_METHOD, serde_json::json!({}))?;
Ok(client)
}
pub async fn start_process(
pub async fn initialize(
&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),
},
);
params: InitializeParams,
) -> Result<InitializeResponse, ExecServerError> {
self.send_request(INITIALIZE_METHOD, params).await
}
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;
}
}
});
async fn send_request<P, R>(&self, method: &str, params: P) -> Result<R, ExecServerError>
where
P: Serialize,
R: DeserializeOwned,
{
let id = RequestId::Integer(self.inner.next_request_id.fetch_add(1, Ordering::SeqCst));
let params = serde_json::to_value(params)?;
let (tx, rx) = oneshot::channel();
self.inner.pending.lock().await.insert(id.clone(), tx);
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));
if let Err(err) = self
.inner
.write_tx
.send(JSONRPCMessage::Request(JSONRPCRequest {
id: id.clone(),
method: method.to_string(),
params: Some(params),
trace: None,
}))
{
let _ = self.inner.pending.lock().await.remove(&id);
return Err(ExecServerError::Protocol(format!(
"failed to queue exec-server request: {err}"
)));
}
Ok(ExecServerProcess {
process_id,
output_rx,
writer_tx,
status,
client: self.clone(),
})
let result = rx.await.map_err(|_| ExecServerError::Closed)??;
Ok(serde_json::from_value(result)?)
}
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> {
fn send_notification<P>(&self, method: &str, params: P) -> Result<(), ExecServerError>
where
P: Serialize,
{
let params = serde_json::to_value(params)?;
self.inner
.write_tx
@@ -372,49 +214,25 @@ impl ExecServerClient {
method: method.to_string(),
params: Some(params),
}))
.map_err(|_| ExecServerError::Closed)
.map_err(|err| {
ExecServerError::Protocol(format!(
"failed to queue exec-server notification: {err}"
))
})
}
}
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,
}),
impl From<JSONRPCErrorError> for ExecServerError {
fn from(error: JSONRPCErrorError) -> Self {
Self::Server {
code: error.code,
message: error.message,
}
}
}
async fn handle_server_message(
inner: &Arc<Inner>,
inner: &Inner,
message: JSONRPCMessage,
) -> Result<(), ExecServerError> {
match message {
@@ -422,76 +240,37 @@ async fn handle_server_message(
if let Some(tx) = inner.pending.lock().await.remove(&id) {
let _ = tx.send(Ok(result));
}
Ok(())
}
JSONRPCMessage::Error(JSONRPCError { id, error }) => {
if let Some(tx) = inner.pending.lock().await.remove(&id) {
let _ = tx.send(Err(error));
Ok(())
} else {
Err(ExecServerError::Server {
code: error.code,
message: error.message,
})
}
}
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
)));
}
JSONRPCMessage::Notification(notification) => Err(ExecServerError::Protocol(format!(
"unexpected exec-server notification: {}",
notification.method
))),
JSONRPCMessage::Request(request) => Err(ExecServerError::Protocol(format!(
"unexpected exec-server request: {}",
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 {
async fn fail_pending_requests(inner: &Inner) {
let mut pending = inner.pending.lock().await;
for (_, tx) in pending.drain() {
let _ = tx.send(Err(JSONRPCErrorError {
code: -32000,
data: None,
message: "exec-server transport closed".to_string(),
data: None,
}));
}
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);
}
}

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

@@ -1,20 +1,11 @@
mod client;
mod protocol;
mod server;
mod server_process;
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;
pub use server_process::ExecServerLaunchCommand;

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

@@ -1,36 +1,10 @@
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 {
@@ -42,102 +16,3 @@ pub struct InitializeParams {
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 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)
}
}

46
codex-rs/exec-server/src/server_process.rs Normal file
View File

@@ -0,0 +1,46 @@
use std::path::PathBuf;
use std::process::Stdio;
use tokio::process::Child;
use tokio::process::ChildStdin;
use tokio::process::ChildStdout;
use tokio::process::Command;
use crate::client::ExecServerError;
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ExecServerLaunchCommand {
pub program: PathBuf,
pub args: Vec<String>,
}
pub(crate) struct SpawnedStdioExecServer {
pub(crate) child: Child,
pub(crate) stdin: ChildStdin,
pub(crate) stdout: ChildStdout,
}
pub(crate) fn spawn_stdio_exec_server(
command: ExecServerLaunchCommand,
) -> Result<SpawnedStdioExecServer, 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())
})?;
Ok(SpawnedStdioExecServer {
child,
stdin,
stdout,
})
}