openai/codex
1
0
Fork 0
mirror of https://github.com/openai/codex.git synced 2026-09-14 11:57:03 +00:00
Code Issues Packages Projects Releases Wiki Activity

exec-server: simplify transport and dispatch shape

Browse Source 
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
starr-openai
2026-03-18 23:05:11 +00:00
parent 52dd39bc95
commit 84bb212aef
9 changed files with 664 additions and 1665 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -31,8 +31,6 @@ use codex_app_server_protocol::JSONRPCResponse;
use codex_app_server_protocol::RequestId;
use serde::Serialize;
use serde_json::Value;
use tokio::io::AsyncRead;
use tokio::io::AsyncWrite;
use tokio::sync::Mutex;
use tokio::sync::broadcast;
use tokio::sync::mpsc;
@@ -76,13 +74,13 @@ use crate::protocol::TerminateResponse;
use crate::protocol::WriteParams;
use crate::protocol::WriteResponse;
use crate::server::ExecServerHandler;
use crate::server::ExecServerOutboundMessage;
use crate::server::ExecServerServerNotification;
impl Default for ExecServerClientConnectOptions {
fn default() -> Self {
Self {
client_name: "codex-core".to_string(),
auth_token: None,
initialize_timeout: INITIALIZE_TIMEOUT,
}
}
@@ -92,6 +90,7 @@ impl From<RemoteExecServerConnectArgs> for ExecServerClientConnectOptions {
fn from(value: RemoteExecServerConnectArgs) -> Self {
Self {
client_name: value.client_name,
auth_token: value.auth_token,
initialize_timeout: value.initialize_timeout,
}
}
@@ -105,6 +104,7 @@ impl RemoteExecServerConnectArgs {
Self {
websocket_url,
client_name,
auth_token: None,
connect_timeout: CONNECT_TIMEOUT,
initialize_timeout: INITIALIZE_TIMEOUT,
}
@@ -123,7 +123,6 @@ struct ExecServerProcess {
process_id: String,
output_rx: broadcast::Receiver<ExecServerOutput>,
status: Arc<RemoteProcessStatus>,
client: ExecServerClient,
}
#[cfg(test)]
@@ -135,18 +134,6 @@ impl ExecServerProcess {
fn has_exited(&self) -> bool {
self.status.has_exited()
}
fn exit_code(&self) -> Option<i32> {
self.status.exit_code()
}
fn terminate(&self) {
let client = self.client.clone();
let process_id = self.process_id.clone();
tokio::spawn(async move {
let _ = client.terminate_session(&process_id).await;
});
}
}
#[cfg(test)]
@@ -168,10 +155,6 @@ impl RemoteProcessStatus {
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() {
@@ -348,21 +331,16 @@ impl ExecServerClient {
pub async fn connect_in_process(
options: ExecServerClientConnectOptions,
) -> Result<Self, ExecServerError> {
let (outbound_tx, mut outgoing_rx) = mpsc::channel::<ExecServerOutboundMessage>(256);
let handler = Arc::new(Mutex::new(ExecServerHandler::new(outbound_tx)));
let (notification_tx, mut notification_rx) =
mpsc::channel::<ExecServerServerNotification>(256);
let handler = Arc::new(Mutex::new(ExecServerHandler::new(notification_tx, None)));
let inner = Arc::new_cyclic(|weak| {
let weak = weak.clone();
let reader_task = tokio::spawn(async move {
while let Some(message) = outgoing_rx.recv().await {
if let Some(inner) = weak.upgrade()
&& let Err(err) = handle_in_process_outbound_message(&inner, message).await
{
warn!(
"in-process exec-server client closing after unexpected response: {err}"
);
handle_transport_shutdown(&inner).await;
return;
while let Some(notification) = notification_rx.recv().await {
if let Some(inner) = weak.upgrade() {
handle_in_process_notification(&inner, notification).await;
}
}
@@ -386,22 +364,6 @@ impl ExecServerClient {
Ok(client)
}
pub async fn connect_stdio<R, W>(
stdin: W,
stdout: R,
options: ExecServerClientConnectOptions,
) -> Result<Self, ExecServerError>
where
R: AsyncRead + Unpin + Send + 'static,
W: AsyncWrite + Unpin + Send + 'static,
{
Self::connect(
JsonRpcConnection::from_stdio(stdout, stdin, "exec-server stdio".to_string()),
options,
)
.await
}
pub async fn connect_websocket(
args: RemoteExecServerConnectArgs,
) -> Result<Self, ExecServerError> {
@@ -521,7 +483,6 @@ impl ExecServerClient {
process_id,
output_rx,
status,
client: self.clone(),
})
}
@@ -648,11 +609,15 @@ impl ExecServerClient {
) -> Result<(), ExecServerError> {
let ExecServerClientConnectOptions {
client_name,
auth_token,
initialize_timeout,
} = options;
timeout(initialize_timeout, async {
let _: InitializeResponse = self
.request_initialize(InitializeParams { client_name })
.request_initialize(InitializeParams {
client_name,
auth_token,
})
.await?;
self.notify(INITIALIZED_METHOD, &serde_json::json!({}))
.await
@@ -735,7 +700,7 @@ impl ExecServerClient {
params: InitializeParams,
) -> Result<InitializeResponse, ExecServerError> {
if let ClientBackend::InProcess { handler } = &self.inner.backend {
return server_result_to_client(handler.lock().await.initialize());
return server_result_to_client(handler.lock().await.initialize(params));
}
self.send_pending_request(INITIALIZE_METHOD, &params, PendingRequest::Initialize)
@@ -825,24 +790,6 @@ async fn send_jsonrpc_request<P: Serialize>(
.map_err(|_| ExecServerError::Closed)
}
async fn handle_in_process_outbound_message(
inner: &Arc<Inner>,
message: ExecServerOutboundMessage,
) -> Result<(), ExecServerError> {
match message {
ExecServerOutboundMessage::Response { .. } | ExecServerOutboundMessage::Error { .. } => {
return Err(ExecServerError::Protocol(
"unexpected in-process RPC response".to_string(),
));
}
ExecServerOutboundMessage::Notification(notification) => {
handle_in_process_notification(inner, notification).await;
}
}
Ok(())
}
async fn handle_in_process_notification(
inner: &Arc<Inner>,
notification: ExecServerServerNotification,

26
codex-rs/exec-server/src/server.rs
View File

@@ -1,21 +1,37 @@
mod filesystem;
mod handler;
mod jsonrpc;
mod processor;
mod routing;
mod transport;
pub(crate) use handler::ExecServerHandler;
pub(crate) use routing::ExecServerOutboundMessage;
pub(crate) use routing::ExecServerServerNotification;
pub(crate) use handler::ExecServerServerNotification;
pub(crate) use jsonrpc::internal_error;
pub(crate) use jsonrpc::invalid_params;
pub(crate) use jsonrpc::invalid_request;
pub(crate) use jsonrpc::unauthorized;
pub use transport::ExecServerTransport;
pub use transport::ExecServerTransportParseError;
#[derive(Debug, Clone, Default, PartialEq, Eq)]
pub struct ExecServerConfig {
pub auth_token: Option<String>,
}
pub async fn run_main() -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
run_main_with_transport(ExecServerTransport::Stdio).await
run_main_with_transport_and_config(ExecServerTransport::default(), ExecServerConfig::default())
.await
}
pub async fn run_main_with_transport(
transport: ExecServerTransport,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
transport::run_transport(transport).await
run_main_with_transport_and_config(transport, ExecServerConfig::default()).await
}
pub async fn run_main_with_transport_and_config(
transport: ExecServerTransport,
config: ExecServerConfig,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
transport::run_transport(transport, config).await
}

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

@@ -25,8 +25,8 @@ use codex_environment::Environment;
use codex_environment::ExecutorFileSystem;
use codex_environment::RemoveOptions;
use crate::server::routing::internal_error;
use crate::server::routing::invalid_request;
use crate::server::internal_error;
use crate::server::invalid_request;
#[derive(Clone)]
pub(crate) struct ExecServerFileSystem {

198
codex-rs/exec-server/src/server/handler.rs
View File

@@ -36,14 +36,19 @@ use crate::protocol::ReadResponse;
use crate::protocol::TerminateResponse;
use crate::protocol::WriteResponse;
use crate::server::filesystem::ExecServerFileSystem;
use crate::server::routing::ExecServerOutboundMessage;
use crate::server::routing::ExecServerServerNotification;
use crate::server::routing::internal_error;
use crate::server::routing::invalid_params;
use crate::server::routing::invalid_request;
use crate::server::internal_error;
use crate::server::invalid_params;
use crate::server::invalid_request;
use crate::server::unauthorized;
const RETAINED_OUTPUT_BYTES_PER_PROCESS: usize = 1024 * 1024;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ExecServerServerNotification {
OutputDelta(ExecOutputDeltaNotification),
Exited(ExecExitedNotification),
}
#[derive(Clone)]
struct RetainedOutputChunk {
seq: u64,
@@ -62,8 +67,9 @@ struct RunningProcess {
}
pub(crate) struct ExecServerHandler {
outbound_tx: mpsc::Sender<ExecServerOutboundMessage>,
notification_tx: mpsc::Sender<ExecServerServerNotification>,
file_system: ExecServerFileSystem,
required_auth_token: Option<String>,
// Keyed by client-chosen logical `processId` scoped to this connection.
// This is a protocol handle, not an OS pid.
processes: Arc<Mutex<HashMap<String, RunningProcess>>>,
@@ -72,10 +78,14 @@ pub(crate) struct ExecServerHandler {
}
impl ExecServerHandler {
pub(crate) fn new(outbound_tx: mpsc::Sender<ExecServerOutboundMessage>) -> Self {
pub(crate) fn new(
notification_tx: mpsc::Sender<ExecServerServerNotification>,
required_auth_token: Option<String>,
) -> Self {
Self {
outbound_tx,
notification_tx,
file_system: ExecServerFileSystem::default(),
required_auth_token,
processes: Arc::new(Mutex::new(HashMap::new())),
initialize_requested: false,
initialized: false,
@@ -105,12 +115,18 @@ impl ExecServerHandler {
pub(crate) fn initialize(
&mut self,
params: crate::protocol::InitializeParams,
) -> Result<InitializeResponse, codex_app_server_protocol::JSONRPCErrorError> {
if self.initialize_requested {
return Err(invalid_request(
"initialize may only be sent once per connection".to_string(),
));
}
if let Some(required_auth_token) = &self.required_auth_token
&& params.auth_token.as_deref() != Some(required_auth_token.as_str())
{
return Err(unauthorized("invalid exec-server auth token".to_string()));
}
self.initialize_requested = true;
Ok(InitializeResponse {
protocol_version: PROTOCOL_VERSION.to_string(),
@@ -210,7 +226,7 @@ impl ExecServerHandler {
ExecOutputStream::Stdout
},
spawned.stdout_rx,
self.outbound_tx.clone(),
self.notification_tx.clone(),
Arc::clone(&self.processes),
Arc::clone(&output_notify),
));
@@ -222,14 +238,14 @@ impl ExecServerHandler {
ExecOutputStream::Stderr
},
spawned.stderr_rx,
self.outbound_tx.clone(),
self.notification_tx.clone(),
Arc::clone(&self.processes),
Arc::clone(&output_notify),
));
tokio::spawn(watch_exit(
process_id.clone(),
spawned.exit_rx,
self.outbound_tx.clone(),
self.notification_tx.clone(),
Arc::clone(&self.processes),
output_notify,
));
@@ -402,153 +418,11 @@ impl ExecServerHandler {
}
}
#[cfg(test)]
impl ExecServerHandler {
async fn handle_message(
&mut self,
message: crate::server::routing::ExecServerInboundMessage,
) -> Result<(), String> {
match message {
crate::server::routing::ExecServerInboundMessage::Request(request) => {
self.handle_request(request).await
}
crate::server::routing::ExecServerInboundMessage::Notification(
crate::server::routing::ExecServerClientNotification::Initialized,
) => self.initialized(),
}
}
async fn handle_request(
&mut self,
request: crate::server::routing::ExecServerRequest,
) -> Result<(), String> {
let outbound = match request {
crate::server::routing::ExecServerRequest::Initialize { request_id, .. } => {
Self::request_outbound(
request_id,
self.initialize()
.map(crate::server::routing::ExecServerResponseMessage::Initialize),
)
}
crate::server::routing::ExecServerRequest::Exec { request_id, params } => {
Self::request_outbound(
request_id,
self.exec(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::Exec),
)
}
crate::server::routing::ExecServerRequest::Read { request_id, params } => {
Self::request_outbound(
request_id,
self.read(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::Read),
)
}
crate::server::routing::ExecServerRequest::Write { request_id, params } => {
Self::request_outbound(
request_id,
self.write(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::Write),
)
}
crate::server::routing::ExecServerRequest::Terminate { request_id, params } => {
Self::request_outbound(
request_id,
self.terminate(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::Terminate),
)
}
crate::server::routing::ExecServerRequest::FsReadFile { request_id, params } => {
Self::request_outbound(
request_id,
self.fs_read_file(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::FsReadFile),
)
}
crate::server::routing::ExecServerRequest::FsWriteFile { request_id, params } => {
Self::request_outbound(
request_id,
self.fs_write_file(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::FsWriteFile),
)
}
crate::server::routing::ExecServerRequest::FsCreateDirectory { request_id, params } => {
Self::request_outbound(
request_id,
self.fs_create_directory(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::FsCreateDirectory),
)
}
crate::server::routing::ExecServerRequest::FsGetMetadata { request_id, params } => {
Self::request_outbound(
request_id,
self.fs_get_metadata(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::FsGetMetadata),
)
}
crate::server::routing::ExecServerRequest::FsReadDirectory { request_id, params } => {
Self::request_outbound(
request_id,
self.fs_read_directory(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::FsReadDirectory),
)
}
crate::server::routing::ExecServerRequest::FsRemove { request_id, params } => {
Self::request_outbound(
request_id,
self.fs_remove(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::FsRemove),
)
}
crate::server::routing::ExecServerRequest::FsCopy { request_id, params } => {
Self::request_outbound(
request_id,
self.fs_copy(params)
.await
.map(crate::server::routing::ExecServerResponseMessage::FsCopy),
)
}
};
self.outbound_tx
.send(outbound)
.await
.map_err(|_| "outbound channel closed".to_string())
}
fn request_outbound(
request_id: codex_app_server_protocol::RequestId,
result: Result<
crate::server::routing::ExecServerResponseMessage,
codex_app_server_protocol::JSONRPCErrorError,
>,
) -> crate::server::routing::ExecServerOutboundMessage {
match result {
Ok(response) => crate::server::routing::ExecServerOutboundMessage::Response {
request_id,
response,
},
Err(error) => {
crate::server::routing::ExecServerOutboundMessage::Error { request_id, error }
}
}
}
}
async fn stream_output(
process_id: String,
stream: ExecOutputStream,
mut receiver: tokio::sync::mpsc::Receiver<Vec<u8>>,
outbound_tx: mpsc::Sender<ExecServerOutboundMessage>,
notification_tx: mpsc::Sender<ExecServerServerNotification>,
processes: Arc<Mutex<HashMap<String, RunningProcess>>>,
output_notify: Arc<Notify>,
) {
@@ -583,10 +457,8 @@ async fn stream_output(
};
output_notify.notify_waiters();
if outbound_tx
.send(ExecServerOutboundMessage::Notification(
ExecServerServerNotification::OutputDelta(notification),
))
if notification_tx
.send(ExecServerServerNotification::OutputDelta(notification))
.await
.is_err()
{
@@ -598,7 +470,7 @@ async fn stream_output(
async fn watch_exit(
process_id: String,
exit_rx: tokio::sync::oneshot::Receiver<i32>,
outbound_tx: mpsc::Sender<ExecServerOutboundMessage>,
notification_tx: mpsc::Sender<ExecServerServerNotification>,
processes: Arc<Mutex<HashMap<String, RunningProcess>>>,
output_notify: Arc<Notify>,
) {
@@ -610,12 +482,12 @@ async fn watch_exit(
}
}
output_notify.notify_waiters();
let _ = outbound_tx
.send(ExecServerOutboundMessage::Notification(
ExecServerServerNotification::Exited(ExecExitedNotification {
let _ = notification_tx
.send(ExecServerServerNotification::Exited(
ExecExitedNotification {
process_id,
exit_code,
}),
},
))
.await;
}

852
codex-rs/exec-server/src/server/handler/tests.rs
View File

@@ -1,322 +1,19 @@
use std::collections::HashMap;
use std::collections::VecDeque;
use std::sync::Arc;
use std::time::Duration;
use pretty_assertions::assert_eq;
use tokio::sync::Notify;
use tokio::time::timeout;
use super::ExecServerHandler;
use super::RetainedOutputChunk;
use super::RunningProcess;
use crate::protocol::ExecOutputStream;
use crate::protocol::ExecParams;
use crate::protocol::ExecSandboxConfig;
use crate::protocol::ExecSandboxMode;
use crate::protocol::InitializeParams;
use crate::protocol::InitializeResponse;
use crate::protocol::PROTOCOL_VERSION;
use crate::protocol::ReadParams;
use crate::protocol::TerminateResponse;
use crate::protocol::TerminateParams;
use crate::protocol::WriteParams;
use crate::server::routing::ExecServerClientNotification;
use crate::server::routing::ExecServerInboundMessage;
use crate::server::routing::ExecServerOutboundMessage;
use crate::server::routing::ExecServerRequest;
use crate::server::routing::ExecServerResponseMessage;
use codex_app_server_protocol::RequestId;
async fn recv_outbound(
outgoing_rx: &mut tokio::sync::mpsc::Receiver<ExecServerOutboundMessage>,
) -> ExecServerOutboundMessage {
let recv_result = timeout(Duration::from_secs(1), outgoing_rx.recv()).await;
let maybe_message = match recv_result {
Ok(maybe_message) => maybe_message,
Err(err) => panic!("timed out waiting for handler output: {err}"),
};
match maybe_message {
Some(message) => message,
None => panic!("handler output channel closed unexpectedly"),
}
}
#[tokio::test]
async fn initialize_response_reports_protocol_version() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(outgoing_tx);
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
.await
{
panic!("initialize should succeed: {err}");
}
assert_eq!(
recv_outbound(&mut outgoing_rx).await,
ExecServerOutboundMessage::Response {
request_id: RequestId::Integer(1),
response: ExecServerResponseMessage::Initialize(InitializeResponse {
protocol_version: PROTOCOL_VERSION.to_string(),
}),
}
);
}
#[tokio::test]
async fn exec_methods_require_initialize() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(outgoing_tx);
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec {
request_id: RequestId::Integer(7),
params: crate::protocol::ExecParams {
process_id: "proc-1".to_string(),
argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()],
cwd: std::env::current_dir().expect("cwd"),
env: HashMap::new(),
tty: true,
arg0: None,
sandbox: None,
},
}))
.await
{
panic!("request handling should not fail the handler: {err}");
}
let ExecServerOutboundMessage::Error { request_id, error } =
recv_outbound(&mut outgoing_rx).await
else {
panic!("expected invalid-request error");
};
assert_eq!(request_id, RequestId::Integer(7));
assert_eq!(error.code, -32600);
assert_eq!(
error.message,
"client must call initialize before using exec methods"
);
}
#[tokio::test]
async fn exec_methods_require_initialized_notification_after_initialize() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2);
let mut handler = ExecServerHandler::new(outgoing_tx);
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
.await
{
panic!("initialize should succeed: {err}");
}
let _ = recv_outbound(&mut outgoing_rx).await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec {
request_id: RequestId::Integer(2),
params: crate::protocol::ExecParams {
process_id: "proc-1".to_string(),
argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()],
cwd: std::env::current_dir().expect("cwd"),
env: HashMap::new(),
tty: true,
arg0: None,
sandbox: None,
},
}))
.await
{
panic!("request handling should not fail the handler: {err}");
}
let ExecServerOutboundMessage::Error { request_id, error } =
recv_outbound(&mut outgoing_rx).await
else {
panic!("expected invalid-request error");
};
assert_eq!(request_id, RequestId::Integer(2));
assert_eq!(error.code, -32600);
assert_eq!(
error.message,
"client must send initialized before using exec methods"
);
}
#[tokio::test]
async fn initialized_before_initialize_is_a_protocol_error() {
let (outgoing_tx, _outgoing_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(outgoing_tx);
let result = handler
.handle_message(ExecServerInboundMessage::Notification(
ExecServerClientNotification::Initialized,
))
.await;
match result {
Err(err) => {
assert_eq!(
err,
"received `initialized` notification before `initialize`"
);
}
Ok(()) => panic!("expected protocol error for early initialized notification"),
}
}
#[tokio::test]
async fn initialize_may_only_be_sent_once_per_connection() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2);
let mut handler = ExecServerHandler::new(outgoing_tx);
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
.await
{
panic!("initialize should succeed: {err}");
}
let _ = recv_outbound(&mut outgoing_rx).await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(2),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
.await
{
panic!("duplicate initialize should not fail the handler: {err}");
}
let ExecServerOutboundMessage::Error { request_id, error } =
recv_outbound(&mut outgoing_rx).await
else {
panic!("expected invalid-request error");
};
assert_eq!(request_id, RequestId::Integer(2));
assert_eq!(error.code, -32600);
assert_eq!(
error.message,
"initialize may only be sent once per connection"
);
}
#[tokio::test]
async fn host_default_sandbox_requests_are_rejected_until_supported() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(3);
let mut handler = ExecServerHandler::new(outgoing_tx);
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
.await
{
panic!("initialize should succeed: {err}");
}
let _ = recv_outbound(&mut outgoing_rx).await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Notification(
ExecServerClientNotification::Initialized,
))
.await
{
panic!("initialized should succeed: {err}");
}
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec {
request_id: RequestId::Integer(2),
params: crate::protocol::ExecParams {
process_id: "proc-1".to_string(),
argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()],
cwd: std::env::current_dir().expect("cwd"),
env: HashMap::new(),
tty: false,
arg0: None,
sandbox: Some(ExecSandboxConfig {
mode: ExecSandboxMode::HostDefault,
}),
},
}))
.await
{
panic!("request handling should not fail the handler: {err}");
}
let ExecServerOutboundMessage::Error { request_id, error } =
recv_outbound(&mut outgoing_rx).await
else {
panic!("expected unsupported sandbox error");
};
assert_eq!(request_id, RequestId::Integer(2));
assert_eq!(error.code, -32600);
assert_eq!(
error.message,
"sandbox mode `hostDefault` is not supported by exec-server yet"
);
}
#[tokio::test]
async fn exec_echoes_client_process_ids() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(4);
let mut handler = ExecServerHandler::new(outgoing_tx);
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
.await
{
panic!("initialize should succeed: {err}");
}
let _ = recv_outbound(&mut outgoing_rx).await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Notification(
ExecServerClientNotification::Initialized,
))
.await
{
panic!("initialized should succeed: {err}");
}
let params = crate::protocol::ExecParams {
process_id: "proc-1".to_string(),
fn exec_params(process_id: &str) -> ExecParams {
ExecParams {
process_id: process_id.to_string(),
argv: vec![
"bash".to_string(),
"-lc".to_string(),
@@ -327,398 +24,203 @@ async fn exec_echoes_client_process_ids() {
tty: false,
arg0: None,
sandbox: None,
};
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec {
request_id: RequestId::Integer(2),
params: params.clone(),
}))
.await
{
panic!("first exec should succeed: {err}");
}
let ExecServerOutboundMessage::Response {
request_id,
response: ExecServerResponseMessage::Exec(first_exec),
} = recv_outbound(&mut outgoing_rx).await
else {
panic!("expected first exec response");
};
assert_eq!(request_id, RequestId::Integer(2));
assert_eq!(first_exec.process_id, "proc-1");
}
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec {
request_id: RequestId::Integer(3),
params: crate::protocol::ExecParams {
process_id: "proc-2".to_string(),
argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()],
..params
},
}))
.await
{
panic!("second exec should succeed: {err}");
}
let ExecServerOutboundMessage::Response {
request_id,
response: ExecServerResponseMessage::Exec(second_exec),
} = recv_outbound(&mut outgoing_rx).await
else {
panic!("expected second exec response");
};
assert_eq!(request_id, RequestId::Integer(3));
assert_eq!(second_exec.process_id, "proc-2");
handler.shutdown().await;
async fn initialized_handler() -> ExecServerHandler {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(8);
let mut handler = ExecServerHandler::new(notification_tx, None);
let response = handler
.initialize(InitializeParams {
client_name: "test".to_string(),
auth_token: None,
})
.expect("initialize should succeed");
assert_eq!(response.protocol_version, PROTOCOL_VERSION);
handler
.initialized()
.expect("initialized notification should succeed");
handler
}
#[tokio::test]
async fn writes_to_pipe_backed_processes_are_rejected() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(4);
let mut handler = ExecServerHandler::new(outgoing_tx);
async fn initialize_reports_protocol_version() {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(notification_tx, None);
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
.await
{
panic!("initialize should succeed: {err}");
}
let _ = recv_outbound(&mut outgoing_rx).await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Notification(
ExecServerClientNotification::Initialized,
))
.await
{
panic!("initialized should succeed: {err}");
}
let response = handler
.initialize(InitializeParams {
client_name: "test".to_string(),
auth_token: None,
})
.expect("initialize should succeed");
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(ExecServerRequest::Exec {
request_id: RequestId::Integer(2),
params: crate::protocol::ExecParams {
process_id: "proc-1".to_string(),
argv: vec![
"bash".to_string(),
"-lc".to_string(),
"sleep 30".to_string(),
],
cwd: std::env::current_dir().expect("cwd"),
env: HashMap::new(),
tty: false,
arg0: None,
sandbox: None,
},
}))
.await
{
panic!("exec should succeed: {err}");
}
let ExecServerOutboundMessage::Response {
response: ExecServerResponseMessage::Exec(exec_response),
..
} = recv_outbound(&mut outgoing_rx).await
else {
panic!("expected exec response");
};
assert_eq!(response.protocol_version, PROTOCOL_VERSION);
}
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Write {
request_id: RequestId::Integer(3),
params: WriteParams {
process_id: exec_response.process_id,
chunk: b"hello\n".to_vec().into(),
},
},
))
.await
{
panic!("write should not fail the handler: {err}");
}
#[tokio::test]
async fn exec_methods_require_initialize() {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1);
let handler = ExecServerHandler::new(notification_tx, None);
let error = handler
.exec(exec_params("proc-1"))
.await
.expect_err("exec should fail before initialize");
let ExecServerOutboundMessage::Error { request_id, error } =
recv_outbound(&mut outgoing_rx).await
else {
panic!("expected stdin-closed error");
};
assert_eq!(request_id, RequestId::Integer(3));
assert_eq!(error.code, -32600);
assert_eq!(error.message, "stdin is closed for process proc-1");
assert_eq!(
error.message,
"client must call initialize before using exec methods"
);
}
#[tokio::test]
async fn exec_methods_require_initialized_notification_after_initialize() {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(notification_tx, None);
let _ = handler
.initialize(InitializeParams {
client_name: "test".to_string(),
auth_token: None,
})
.expect("initialize should succeed");
let error = handler
.exec(exec_params("proc-1"))
.await
.expect_err("exec should fail before initialized notification");
assert_eq!(error.code, -32600);
assert_eq!(
error.message,
"client must send initialized before using exec methods"
);
}
#[tokio::test]
async fn initialized_before_initialize_is_a_protocol_error() {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(notification_tx, None);
let error = handler
.initialized()
.expect_err("expected protocol error for early initialized notification");
assert_eq!(
error,
"received `initialized` notification before `initialize`"
);
}
#[tokio::test]
async fn initialize_may_only_be_sent_once_per_connection() {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(notification_tx, None);
let _ = handler
.initialize(InitializeParams {
client_name: "test".to_string(),
auth_token: None,
})
.expect("first initialize should succeed");
let error = handler
.initialize(InitializeParams {
client_name: "test".to_string(),
auth_token: None,
})
.expect_err("duplicate initialize should fail");
assert_eq!(error.code, -32600);
assert_eq!(
error.message,
"initialize may only be sent once per connection"
);
}
#[tokio::test]
async fn initialize_rejects_invalid_auth_token() {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(notification_tx, Some("secret-token".to_string()));
let error = handler
.initialize(InitializeParams {
client_name: "test".to_string(),
auth_token: Some("wrong-token".to_string()),
})
.expect_err("invalid auth token should fail");
assert_eq!(error.code, -32001);
assert_eq!(error.message, "invalid exec-server auth token");
}
#[tokio::test]
async fn exec_rejects_host_default_sandbox_mode() {
let handler = initialized_handler().await;
let error = handler
.exec(ExecParams {
sandbox: Some(ExecSandboxConfig {
mode: ExecSandboxMode::HostDefault,
}),
..exec_params("proc-1")
})
.await
.expect_err("hostDefault sandbox should be rejected");
assert_eq!(error.code, -32600);
assert_eq!(
error.message,
"sandbox mode `hostDefault` is not supported by exec-server yet"
);
}
#[tokio::test]
async fn exec_rejects_duplicate_process_ids() {
let handler = initialized_handler().await;
let first = handler
.exec(exec_params("proc-1"))
.await
.expect("first exec should succeed");
assert_eq!(first.process_id, "proc-1");
let error = handler
.exec(exec_params("proc-1"))
.await
.expect_err("duplicate process id should fail");
assert_eq!(error.code, -32600);
assert_eq!(error.message, "process proc-1 already exists");
handler.shutdown().await;
}
#[tokio::test]
async fn writes_to_unknown_processes_are_rejected() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2);
let mut handler = ExecServerHandler::new(outgoing_tx);
async fn write_rejects_unknown_process_ids() {
let handler = initialized_handler().await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
let error = handler
.write(WriteParams {
process_id: "missing".to_string(),
chunk: b"input".to_vec().into(),
})
.await
{
panic!("initialize should succeed: {err}");
}
let _ = recv_outbound(&mut outgoing_rx).await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Notification(
ExecServerClientNotification::Initialized,
))
.await
{
panic!("initialized should succeed: {err}");
}
.expect_err("writing to an unknown process should fail");
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Write {
request_id: RequestId::Integer(2),
params: WriteParams {
process_id: "missing".to_string(),
chunk: b"hello\n".to_vec().into(),
},
},
))
.await
{
panic!("write should not fail the handler: {err}");
}
let ExecServerOutboundMessage::Error { request_id, error } =
recv_outbound(&mut outgoing_rx).await
else {
panic!("expected unknown-process error");
};
assert_eq!(request_id, RequestId::Integer(2));
assert_eq!(error.code, -32600);
assert_eq!(error.message, "unknown process id missing");
}
#[tokio::test]
async fn terminate_unknown_processes_report_running_false() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2);
let mut handler = ExecServerHandler::new(outgoing_tx);
async fn terminate_reports_missing_processes_as_not_running() {
let handler = initialized_handler().await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
.await
{
panic!("initialize should succeed: {err}");
}
let _ = recv_outbound(&mut outgoing_rx).await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Notification(
ExecServerClientNotification::Initialized,
))
.await
{
panic!("initialized should succeed: {err}");
}
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Terminate {
request_id: RequestId::Integer(2),
params: crate::protocol::TerminateParams {
process_id: "missing".to_string(),
},
},
))
.await
{
panic!("terminate should not fail the handler: {err}");
}
assert_eq!(
recv_outbound(&mut outgoing_rx).await,
ExecServerOutboundMessage::Response {
request_id: RequestId::Integer(2),
response: ExecServerResponseMessage::Terminate(TerminateResponse { running: false }),
}
);
}
#[tokio::test]
async fn terminate_keeps_process_ids_reserved() {
let (outgoing_tx, mut outgoing_rx) = tokio::sync::mpsc::channel(2);
let mut handler = ExecServerHandler::new(outgoing_tx);
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test".to_string(),
},
},
))
.await
{
panic!("initialize should succeed: {err}");
}
let _ = recv_outbound(&mut outgoing_rx).await;
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Notification(
ExecServerClientNotification::Initialized,
))
.await
{
panic!("initialized should succeed: {err}");
}
let spawned = codex_utils_pty::spawn_pipe_process_no_stdin(
"bash",
&["-lc".to_string(), "sleep 30".to_string()],
std::env::current_dir().expect("cwd").as_path(),
&HashMap::new(),
&None,
)
.await
.expect("spawn test process");
{
let mut process_map = handler.processes.lock().await;
process_map.insert(
"proc-1".to_string(),
super::RunningProcess {
session: spawned.session,
tty: false,
output: std::collections::VecDeque::new(),
retained_bytes: 0,
next_seq: 1,
exit_code: None,
output_notify: Arc::new(Notify::new()),
},
);
}
if let Err(err) = handler
.handle_message(ExecServerInboundMessage::Request(
ExecServerRequest::Terminate {
request_id: RequestId::Integer(2),
params: crate::protocol::TerminateParams {
process_id: "proc-1".to_string(),
},
},
))
.await
{
panic!("terminate should not fail the handler: {err}");
}
assert_eq!(
recv_outbound(&mut outgoing_rx).await,
ExecServerOutboundMessage::Response {
request_id: RequestId::Integer(2),
response: ExecServerResponseMessage::Terminate(TerminateResponse { running: true }),
}
);
assert!(
handler.processes.lock().await.contains_key("proc-1"),
"terminated ids should stay reserved until exit cleanup removes them"
);
handler.shutdown().await;
}
#[tokio::test]
async fn read_paginates_retained_output_without_skipping_omitted_chunks() {
let (outgoing_tx, _outgoing_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(outgoing_tx);
let _ = handler.initialize().expect("initialize should succeed");
handler.initialized().expect("initialized should succeed");
let spawned = codex_utils_pty::spawn_pipe_process_no_stdin(
"bash",
&["-lc".to_string(), "true".to_string()],
std::env::current_dir().expect("cwd").as_path(),
&HashMap::new(),
&None,
)
.await
.expect("spawn test process");
{
let mut process_map = handler.processes.lock().await;
process_map.insert(
"proc-1".to_string(),
RunningProcess {
session: spawned.session,
tty: false,
output: VecDeque::from([
RetainedOutputChunk {
seq: 1,
stream: ExecOutputStream::Stdout,
chunk: b"abc".to_vec(),
},
RetainedOutputChunk {
seq: 2,
stream: ExecOutputStream::Stderr,
chunk: b"def".to_vec(),
},
]),
retained_bytes: 6,
next_seq: 3,
exit_code: None,
output_notify: Arc::new(Notify::new()),
},
);
}
let first = handler
.read(ReadParams {
process_id: "proc-1".to_string(),
after_seq: Some(0),
max_bytes: Some(3),
wait_ms: Some(0),
let response = handler
.terminate(TerminateParams {
process_id: "missing".to_string(),
})
.await
.expect("first read should succeed");
.expect("terminate should succeed");
assert_eq!(first.chunks.len(), 1);
assert_eq!(first.chunks[0].seq, 1);
assert_eq!(first.chunks[0].stream, ExecOutputStream::Stdout);
assert_eq!(first.chunks[0].chunk.clone().into_inner(), b"abc".to_vec());
assert_eq!(first.next_seq, 2);
let second = handler
.read(ReadParams {
process_id: "proc-1".to_string(),
after_seq: Some(first.next_seq - 1),
max_bytes: Some(3),
wait_ms: Some(0),
})
.await
.expect("second read should succeed");
assert_eq!(second.chunks.len(), 1);
assert_eq!(second.chunks[0].seq, 2);
assert_eq!(second.chunks[0].stream, ExecOutputStream::Stderr);
assert_eq!(second.chunks[0].chunk.clone().into_inner(), b"def".to_vec());
assert_eq!(second.next_seq, 3);
handler.shutdown().await;
assert_eq!(response.running, false);
}

33
codex-rs/exec-server/src/server/jsonrpc.rs Normal file
View File

@@ -0,0 +1,33 @@
use codex_app_server_protocol::JSONRPCErrorError;
pub(crate) fn invalid_request(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32600,
data: None,
message,
}
}
pub(crate) fn invalid_params(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32602,
data: None,
message,
}
}
pub(crate) fn internal_error(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32603,
data: None,
message,
}
}
pub(crate) fn unauthorized(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32001,
data: None,
message,
}
}

493
codex-rs/exec-server/src/server/processor.rs
View File

@@ -1,3 +1,12 @@
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 tokio::sync::mpsc;
use tracing::debug;
use tracing::warn;
@@ -5,32 +14,49 @@ use tracing::warn;
use crate::connection::CHANNEL_CAPACITY;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
use crate::server::handler::ExecServerHandler;
use crate::server::routing::ExecServerClientNotification;
use crate::server::routing::ExecServerInboundMessage;
use crate::server::routing::ExecServerOutboundMessage;
use crate::server::routing::ExecServerRequest;
use crate::server::routing::ExecServerResponseMessage;
use crate::server::routing::RoutedExecServerMessage;
use crate::server::routing::encode_outbound_message;
use crate::server::routing::route_jsonrpc_message;
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::FS_COPY_METHOD;
use crate::protocol::FS_CREATE_DIRECTORY_METHOD;
use crate::protocol::FS_GET_METADATA_METHOD;
use crate::protocol::FS_READ_DIRECTORY_METHOD;
use crate::protocol::FS_READ_FILE_METHOD;
use crate::protocol::FS_REMOVE_METHOD;
use crate::protocol::FS_WRITE_FILE_METHOD;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::server::ExecServerConfig;
use crate::server::ExecServerHandler;
use crate::server::ExecServerServerNotification;
use crate::server::internal_error;
use crate::server::invalid_params;
use crate::server::invalid_request;
pub(crate) async fn run_connection(connection: JsonRpcConnection) {
pub(crate) async fn run_connection(connection: JsonRpcConnection, config: ExecServerConfig) {
let (json_outgoing_tx, mut incoming_rx, _connection_tasks) = connection.into_parts();
let (outgoing_tx, mut outgoing_rx) =
mpsc::channel::<ExecServerOutboundMessage>(CHANNEL_CAPACITY);
let mut handler = ExecServerHandler::new(outgoing_tx.clone());
let json_outgoing_tx_for_notifications = json_outgoing_tx.clone();
let (notification_tx, mut notification_rx) =
mpsc::channel::<ExecServerServerNotification>(CHANNEL_CAPACITY);
let mut handler = ExecServerHandler::new(notification_tx, config.auth_token);
let outbound_task = tokio::spawn(async move {
while let Some(message) = outgoing_rx.recv().await {
let json_message = match encode_outbound_message(message) {
while let Some(notification) = notification_rx.recv().await {
let json_message = match notification_message(notification) {
Ok(json_message) => json_message,
Err(err) => {
warn!("failed to serialize exec-server outbound message: {err}");
warn!("failed to serialize exec-server notification: {err}");
break;
}
};
if json_outgoing_tx.send(json_message).await.is_err() {
if json_outgoing_tx_for_notifications
.send(json_message)
.await
.is_err()
{
break;
}
}
@@ -38,24 +64,20 @@ pub(crate) async fn run_connection(connection: JsonRpcConnection) {
while let Some(event) = incoming_rx.recv().await {
match event {
JsonRpcConnectionEvent::Message(message) => match route_jsonrpc_message(message) {
Ok(RoutedExecServerMessage::Inbound(message)) => {
if let Err(err) = dispatch_to_handler(&mut handler, message, &outgoing_tx).await
{
JsonRpcConnectionEvent::Message(message) => {
let maybe_response = match handle_connection_message(&mut handler, message).await {
Ok(maybe_response) => maybe_response,
Err(err) => {
warn!("closing exec-server connection after protocol error: {err}");
break;
}
}
Ok(RoutedExecServerMessage::ImmediateOutbound(message)) => {
if outgoing_tx.send(message).await.is_err() {
break;
}
}
Err(err) => {
warn!("closing exec-server connection after protocol error: {err}");
};
if let Some(response) = maybe_response
&& json_outgoing_tx.send(response).await.is_err()
{
break;
}
},
}
JsonRpcConnectionEvent::Disconnected { reason } => {
if let Some(reason) = reason {
debug!("exec-server connection disconnected: {reason}");
@@ -67,122 +89,319 @@ pub(crate) async fn run_connection(connection: JsonRpcConnection) {
handler.shutdown().await;
drop(handler);
drop(outgoing_tx);
let _ = outbound_task.await;
}
async fn dispatch_to_handler(
async fn handle_connection_message(
handler: &mut ExecServerHandler,
message: ExecServerInboundMessage,
outgoing_tx: &mpsc::Sender<ExecServerOutboundMessage>,
) -> Result<(), String> {
message: JSONRPCMessage,
) -> Result<Option<JSONRPCMessage>, String> {
match message {
ExecServerInboundMessage::Request(request) => {
let outbound = match request {
ExecServerRequest::Initialize { request_id, .. } => request_outbound(
request_id,
handler
.initialize()
.map(ExecServerResponseMessage::Initialize),
),
ExecServerRequest::Exec { request_id, params } => request_outbound(
request_id,
handler
.exec(params)
.await
.map(ExecServerResponseMessage::Exec),
),
ExecServerRequest::Read { request_id, params } => request_outbound(
request_id,
handler
.read(params)
.await
.map(ExecServerResponseMessage::Read),
),
ExecServerRequest::Write { request_id, params } => request_outbound(
request_id,
handler
.write(params)
.await
.map(ExecServerResponseMessage::Write),
),
ExecServerRequest::Terminate { request_id, params } => request_outbound(
request_id,
handler
.terminate(params)
.await
.map(ExecServerResponseMessage::Terminate),
),
ExecServerRequest::FsReadFile { request_id, params } => request_outbound(
request_id,
handler
.fs_read_file(params)
.await
.map(ExecServerResponseMessage::FsReadFile),
),
ExecServerRequest::FsWriteFile { request_id, params } => request_outbound(
request_id,
handler
.fs_write_file(params)
.await
.map(ExecServerResponseMessage::FsWriteFile),
),
ExecServerRequest::FsCreateDirectory { request_id, params } => request_outbound(
request_id,
handler
.fs_create_directory(params)
.await
.map(ExecServerResponseMessage::FsCreateDirectory),
),
ExecServerRequest::FsGetMetadata { request_id, params } => request_outbound(
request_id,
handler
.fs_get_metadata(params)
.await
.map(ExecServerResponseMessage::FsGetMetadata),
),
ExecServerRequest::FsReadDirectory { request_id, params } => request_outbound(
request_id,
handler
.fs_read_directory(params)
.await
.map(ExecServerResponseMessage::FsReadDirectory),
),
ExecServerRequest::FsRemove { request_id, params } => request_outbound(
request_id,
handler
.fs_remove(params)
.await
.map(ExecServerResponseMessage::FsRemove),
),
ExecServerRequest::FsCopy { request_id, params } => request_outbound(
request_id,
handler
.fs_copy(params)
.await
.map(ExecServerResponseMessage::FsCopy),
),
};
outgoing_tx
.send(outbound)
.await
.map_err(|_| "outbound channel closed".to_string())
JSONRPCMessage::Request(request) => Ok(Some(dispatch_request(handler, request).await)),
JSONRPCMessage::Notification(notification) => {
handle_notification(handler, notification)?;
Ok(None)
}
ExecServerInboundMessage::Notification(ExecServerClientNotification::Initialized) => {
handler.initialized()
JSONRPCMessage::Response(response) => Err(format!(
"unexpected client response for request id {:?}",
response.id
)),
JSONRPCMessage::Error(error) => Err(format!(
"unexpected client error for request id {:?}",
error.id
)),
}
}
async fn dispatch_request(
handler: &mut ExecServerHandler,
request: JSONRPCRequest,
) -> JSONRPCMessage {
let JSONRPCRequest {
id,
method,
params,
trace: _,
} = request;
let params = params.unwrap_or(serde_json::Value::Null);
match method.as_str() {
INITIALIZE_METHOD => request_response(
id,
parse_params(params).and_then(|params| handler.initialize(params)),
),
EXEC_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.exec(params)).await,
),
EXEC_READ_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.read(params)).await,
),
EXEC_WRITE_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.write(params)).await,
),
EXEC_TERMINATE_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.terminate(params)).await,
),
FS_READ_FILE_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.fs_read_file(params)).await,
),
FS_WRITE_FILE_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.fs_write_file(params)).await,
),
FS_CREATE_DIRECTORY_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.fs_create_directory(params)).await,
),
FS_GET_METADATA_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.fs_get_metadata(params)).await,
),
FS_READ_DIRECTORY_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.fs_read_directory(params)).await,
),
FS_REMOVE_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.fs_remove(params)).await,
),
FS_COPY_METHOD => request_response(
id,
dispatch_async_request(params, |params| handler.fs_copy(params)).await,
),
other => jsonrpc_error_response(id, invalid_request(format!("unknown method: {other}"))),
}
}
fn handle_notification(
handler: &mut ExecServerHandler,
notification: JSONRPCNotification,
) -> Result<(), String> {
match notification.method.as_str() {
INITIALIZED_METHOD => handler.initialized(),
other => Err(format!("unexpected notification method: {other}")),
}
}
fn parse_params<P>(params: serde_json::Value) -> Result<P, JSONRPCErrorError>
where
P: DeserializeOwned,
{
serde_json::from_value(params).map_err(|err| invalid_params(err.to_string()))
}
async fn dispatch_async_request<P, T, F, Fut>(
params: serde_json::Value,
f: F,
) -> Result<T, JSONRPCErrorError>
where
P: DeserializeOwned,
F: FnOnce(P) -> Fut,
Fut: std::future::Future<Output = Result<T, JSONRPCErrorError>>,
{
match parse_params(params) {
Ok(params) => f(params).await,
Err(err) => Err(err),
}
}
fn request_response<T>(
request_id: RequestId,
result: Result<T, JSONRPCErrorError>,
) -> JSONRPCMessage
where
T: Serialize,
{
match result.and_then(serialize_response) {
Ok(result) => JSONRPCMessage::Response(JSONRPCResponse {
id: request_id,
result,
}),
Err(error) => JSONRPCMessage::Error(JSONRPCError {
id: request_id,
error,
}),
}
}
fn serialize_response<T>(response: T) -> Result<serde_json::Value, JSONRPCErrorError>
where
T: Serialize,
{
serde_json::to_value(response).map_err(|err| internal_error(err.to_string()))
}
fn jsonrpc_error_response(request_id: RequestId, error: JSONRPCErrorError) -> JSONRPCMessage {
JSONRPCMessage::Error(JSONRPCError {
id: request_id,
error,
})
}
fn notification_message(
notification: ExecServerServerNotification,
) -> Result<JSONRPCMessage, serde_json::Error> {
match notification {
ExecServerServerNotification::OutputDelta(params) => {
typed_notification(EXEC_OUTPUT_DELTA_METHOD, params)
}
ExecServerServerNotification::Exited(params) => {
typed_notification(EXEC_EXITED_METHOD, params)
}
}
}
fn request_outbound(
request_id: codex_app_server_protocol::RequestId,
result: Result<ExecServerResponseMessage, codex_app_server_protocol::JSONRPCErrorError>,
) -> ExecServerOutboundMessage {
match result {
Ok(response) => ExecServerOutboundMessage::Response {
request_id,
response,
},
Err(error) => ExecServerOutboundMessage::Error { request_id, error },
fn typed_notification<T>(method: &str, params: T) -> Result<JSONRPCMessage, serde_json::Error>
where
T: Serialize,
{
Ok(JSONRPCMessage::Notification(JSONRPCNotification {
method: method.to_string(),
params: Some(serde_json::to_value(params)?),
}))
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use pretty_assertions::assert_eq;
use super::dispatch_request;
use super::handle_connection_message;
use super::notification_message;
use crate::protocol::EXEC_METHOD;
use crate::protocol::EXEC_OUTPUT_DELTA_METHOD;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecOutputStream;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeParams;
use crate::protocol::PROTOCOL_VERSION;
use crate::server::ExecServerHandler;
use crate::server::ExecServerServerNotification;
use codex_app_server_protocol::JSONRPCError;
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;
#[tokio::test]
async fn dispatch_initialize_returns_initialize_response() {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(notification_tx, None);
let message = dispatch_request(
&mut handler,
JSONRPCRequest {
id: RequestId::Integer(1),
method: INITIALIZE_METHOD.to_string(),
params: Some(
serde_json::to_value(InitializeParams {
client_name: "test".to_string(),
auth_token: None,
})
.expect("serialize initialize params"),
),
trace: None,
},
)
.await;
assert_eq!(
message,
JSONRPCMessage::Response(JSONRPCResponse {
id: RequestId::Integer(1),
result: serde_json::json!({
"protocolVersion": PROTOCOL_VERSION,
}),
})
);
}
#[tokio::test]
async fn dispatch_exec_returns_invalid_request_before_initialize() {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(notification_tx, None);
let message = dispatch_request(
&mut handler,
JSONRPCRequest {
id: RequestId::Integer(7),
method: EXEC_METHOD.to_string(),
params: Some(serde_json::json!({
"processId": "proc-1",
"argv": ["bash", "-lc", "true"],
"cwd": std::env::current_dir().expect("cwd"),
"env": HashMap::<String, String>::new(),
"tty": true,
"arg0": null,
"sandbox": null,
})),
trace: None,
},
)
.await;
let JSONRPCMessage::Error(JSONRPCError { id, error }) = message else {
panic!("expected invalid-request error");
};
assert_eq!(id, RequestId::Integer(7));
assert_eq!(error.code, -32600);
assert_eq!(
error.message,
"client must call initialize before using exec methods"
);
}
#[tokio::test]
async fn initialized_notification_before_initialize_is_protocol_error() {
let (notification_tx, _notification_rx) = tokio::sync::mpsc::channel(1);
let mut handler = ExecServerHandler::new(notification_tx, None);
let err = handle_connection_message(
&mut handler,
JSONRPCMessage::Notification(JSONRPCNotification {
method: INITIALIZED_METHOD.to_string(),
params: Some(serde_json::json!({})),
}),
)
.await
.expect_err("expected early initialized to fail");
assert_eq!(
err,
"received `initialized` notification before `initialize`"
);
}
#[test]
fn notification_message_serializes_process_output() {
let message = notification_message(ExecServerServerNotification::OutputDelta(
ExecOutputDeltaNotification {
process_id: "proc-1".to_string(),
stream: ExecOutputStream::Stdout,
chunk: b"hello".to_vec().into(),
},
))
.expect("serialize notification");
assert_eq!(
message,
JSONRPCMessage::Notification(JSONRPCNotification {
method: EXEC_OUTPUT_DELTA_METHOD.to_string(),
params: Some(serde_json::json!({
"processId": "proc-1",
"stream": "stdout",
"chunk": "aGVsbG8=",
})),
})
);
}
}

585
codex-rs/exec-server/src/server/routing.rs
View File

@@ -1,585 +0,0 @@
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::de::DeserializeOwned;
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::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::FS_COPY_METHOD;
use crate::protocol::FS_CREATE_DIRECTORY_METHOD;
use crate::protocol::FS_GET_METADATA_METHOD;
use crate::protocol::FS_READ_DIRECTORY_METHOD;
use crate::protocol::FS_READ_FILE_METHOD;
use crate::protocol::FS_REMOVE_METHOD;
use crate::protocol::FS_WRITE_FILE_METHOD;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeParams;
use crate::protocol::InitializeResponse;
use crate::protocol::ReadParams;
use crate::protocol::ReadResponse;
use crate::protocol::TerminateParams;
use crate::protocol::TerminateResponse;
use crate::protocol::WriteParams;
use crate::protocol::WriteResponse;
use codex_app_server_protocol::FsCopyParams;
use codex_app_server_protocol::FsCopyResponse;
use codex_app_server_protocol::FsCreateDirectoryParams;
use codex_app_server_protocol::FsCreateDirectoryResponse;
use codex_app_server_protocol::FsGetMetadataParams;
use codex_app_server_protocol::FsGetMetadataResponse;
use codex_app_server_protocol::FsReadDirectoryParams;
use codex_app_server_protocol::FsReadDirectoryResponse;
use codex_app_server_protocol::FsReadFileParams;
use codex_app_server_protocol::FsReadFileResponse;
use codex_app_server_protocol::FsRemoveParams;
use codex_app_server_protocol::FsRemoveResponse;
use codex_app_server_protocol::FsWriteFileParams;
use codex_app_server_protocol::FsWriteFileResponse;
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ExecServerInboundMessage {
Request(ExecServerRequest),
Notification(ExecServerClientNotification),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ExecServerRequest {
Initialize {
request_id: RequestId,
params: InitializeParams,
},
Exec {
request_id: RequestId,
params: ExecParams,
},
Read {
request_id: RequestId,
params: ReadParams,
},
Write {
request_id: RequestId,
params: WriteParams,
},
Terminate {
request_id: RequestId,
params: TerminateParams,
},
FsReadFile {
request_id: RequestId,
params: FsReadFileParams,
},
FsWriteFile {
request_id: RequestId,
params: FsWriteFileParams,
},
FsCreateDirectory {
request_id: RequestId,
params: FsCreateDirectoryParams,
},
FsGetMetadata {
request_id: RequestId,
params: FsGetMetadataParams,
},
FsReadDirectory {
request_id: RequestId,
params: FsReadDirectoryParams,
},
FsRemove {
request_id: RequestId,
params: FsRemoveParams,
},
FsCopy {
request_id: RequestId,
params: FsCopyParams,
},
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ExecServerClientNotification {
Initialized,
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum ExecServerOutboundMessage {
Response {
request_id: RequestId,
response: ExecServerResponseMessage,
},
Error {
request_id: RequestId,
error: JSONRPCErrorError,
},
Notification(ExecServerServerNotification),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ExecServerResponseMessage {
Initialize(InitializeResponse),
Exec(ExecResponse),
Read(ReadResponse),
Write(WriteResponse),
Terminate(TerminateResponse),
FsReadFile(FsReadFileResponse),
FsWriteFile(FsWriteFileResponse),
FsCreateDirectory(FsCreateDirectoryResponse),
FsGetMetadata(FsGetMetadataResponse),
FsReadDirectory(FsReadDirectoryResponse),
FsRemove(FsRemoveResponse),
FsCopy(FsCopyResponse),
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub(crate) enum ExecServerServerNotification {
OutputDelta(ExecOutputDeltaNotification),
Exited(ExecExitedNotification),
}
#[derive(Debug, Clone, PartialEq)]
pub(crate) enum RoutedExecServerMessage {
Inbound(ExecServerInboundMessage),
ImmediateOutbound(ExecServerOutboundMessage),
}
pub(crate) fn route_jsonrpc_message(
message: JSONRPCMessage,
) -> Result<RoutedExecServerMessage, String> {
match message {
JSONRPCMessage::Request(request) => route_request(request),
JSONRPCMessage::Notification(notification) => route_notification(notification),
JSONRPCMessage::Response(response) => Err(format!(
"unexpected client response for request id {:?}",
response.id
)),
JSONRPCMessage::Error(error) => Err(format!(
"unexpected client error for request id {:?}",
error.id
)),
}
}
pub(crate) fn encode_outbound_message(
message: ExecServerOutboundMessage,
) -> Result<JSONRPCMessage, serde_json::Error> {
match message {
ExecServerOutboundMessage::Response {
request_id,
response,
} => Ok(JSONRPCMessage::Response(JSONRPCResponse {
id: request_id,
result: serialize_response(response)?,
})),
ExecServerOutboundMessage::Error { request_id, error } => {
Ok(JSONRPCMessage::Error(JSONRPCError {
id: request_id,
error,
}))
}
ExecServerOutboundMessage::Notification(notification) => Ok(JSONRPCMessage::Notification(
serialize_notification(notification)?,
)),
}
}
pub(crate) fn invalid_request(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32600,
data: None,
message,
}
}
pub(crate) fn invalid_params(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32602,
data: None,
message,
}
}
pub(crate) fn internal_error(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32603,
data: None,
message,
}
}
fn route_request(request: JSONRPCRequest) -> Result<RoutedExecServerMessage, String> {
match request.method.as_str() {
INITIALIZE_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::Initialize { request_id, params }
})),
EXEC_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::Exec { request_id, params }
})),
EXEC_READ_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::Read { request_id, params }
})),
EXEC_WRITE_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::Write { request_id, params }
})),
EXEC_TERMINATE_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::Terminate { request_id, params }
})),
FS_READ_FILE_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::FsReadFile { request_id, params }
})),
FS_WRITE_FILE_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::FsWriteFile { request_id, params }
})),
FS_CREATE_DIRECTORY_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::FsCreateDirectory { request_id, params }
})),
FS_GET_METADATA_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::FsGetMetadata { request_id, params }
})),
FS_READ_DIRECTORY_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::FsReadDirectory { request_id, params }
})),
FS_REMOVE_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::FsRemove { request_id, params }
})),
FS_COPY_METHOD => Ok(parse_request_params(request, |request_id, params| {
ExecServerRequest::FsCopy { request_id, params }
})),
other => Ok(RoutedExecServerMessage::ImmediateOutbound(
ExecServerOutboundMessage::Error {
request_id: request.id,
error: invalid_request(format!("unknown method: {other}")),
},
)),
}
}
fn route_notification(
notification: JSONRPCNotification,
) -> Result<RoutedExecServerMessage, String> {
match notification.method.as_str() {
INITIALIZED_METHOD => Ok(RoutedExecServerMessage::Inbound(
ExecServerInboundMessage::Notification(ExecServerClientNotification::Initialized),
)),
other => Err(format!("unexpected notification method: {other}")),
}
}
fn parse_request_params<P, F>(request: JSONRPCRequest, build: F) -> RoutedExecServerMessage
where
P: DeserializeOwned,
F: FnOnce(RequestId, P) -> ExecServerRequest,
{
let request_id = request.id;
match serde_json::from_value::<P>(request.params.unwrap_or(serde_json::Value::Null)) {
Ok(params) => RoutedExecServerMessage::Inbound(ExecServerInboundMessage::Request(build(
request_id, params,
))),
Err(err) => RoutedExecServerMessage::ImmediateOutbound(ExecServerOutboundMessage::Error {
request_id,
error: invalid_params(err.to_string()),
}),
}
}
fn serialize_response(
response: ExecServerResponseMessage,
) -> Result<serde_json::Value, serde_json::Error> {
match response {
ExecServerResponseMessage::Initialize(response) => serde_json::to_value(response),
ExecServerResponseMessage::Exec(response) => serde_json::to_value(response),
ExecServerResponseMessage::Read(response) => serde_json::to_value(response),
ExecServerResponseMessage::Write(response) => serde_json::to_value(response),
ExecServerResponseMessage::Terminate(response) => serde_json::to_value(response),
ExecServerResponseMessage::FsReadFile(response) => serde_json::to_value(response),
ExecServerResponseMessage::FsWriteFile(response) => serde_json::to_value(response),
ExecServerResponseMessage::FsCreateDirectory(response) => serde_json::to_value(response),
ExecServerResponseMessage::FsGetMetadata(response) => serde_json::to_value(response),
ExecServerResponseMessage::FsReadDirectory(response) => serde_json::to_value(response),
ExecServerResponseMessage::FsRemove(response) => serde_json::to_value(response),
ExecServerResponseMessage::FsCopy(response) => serde_json::to_value(response),
}
}
fn serialize_notification(
notification: ExecServerServerNotification,
) -> Result<JSONRPCNotification, serde_json::Error> {
match notification {
ExecServerServerNotification::OutputDelta(params) => Ok(JSONRPCNotification {
method: EXEC_OUTPUT_DELTA_METHOD.to_string(),
params: Some(serde_json::to_value(params)?),
}),
ExecServerServerNotification::Exited(params) => Ok(JSONRPCNotification {
method: EXEC_EXITED_METHOD.to_string(),
params: Some(serde_json::to_value(params)?),
}),
}
}
#[cfg(test)]
mod tests {
use pretty_assertions::assert_eq;
use serde_json::json;
use super::ExecServerClientNotification;
use super::ExecServerInboundMessage;
use super::ExecServerOutboundMessage;
use super::ExecServerRequest;
use super::ExecServerResponseMessage;
use super::ExecServerServerNotification;
use super::RoutedExecServerMessage;
use super::encode_outbound_message;
use super::route_jsonrpc_message;
use crate::protocol::EXEC_EXITED_METHOD;
use crate::protocol::EXEC_METHOD;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::ExecSandboxConfig;
use crate::protocol::ExecSandboxMode;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeParams;
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;
#[test]
fn routes_initialize_requests_to_typed_variants() {
let routed = route_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest {
id: RequestId::Integer(1),
method: INITIALIZE_METHOD.to_string(),
params: Some(json!({ "clientName": "test-client" })),
trace: None,
}))
.expect("initialize request should route");
assert_eq!(
routed,
RoutedExecServerMessage::Inbound(ExecServerInboundMessage::Request(
ExecServerRequest::Initialize {
request_id: RequestId::Integer(1),
params: InitializeParams {
client_name: "test-client".to_string(),
},
},
))
);
}
#[test]
fn malformed_exec_params_return_immediate_error_outbound() {
let routed = route_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest {
id: RequestId::Integer(2),
method: EXEC_METHOD.to_string(),
params: Some(json!({ "processId": "proc-1" })),
trace: None,
}))
.expect("exec request should route");
let RoutedExecServerMessage::ImmediateOutbound(ExecServerOutboundMessage::Error {
request_id,
error,
}) = routed
else {
panic!("expected invalid-params error outbound");
};
assert_eq!(request_id, RequestId::Integer(2));
assert_eq!(error.code, -32602);
}
#[test]
fn routes_initialized_notifications_to_typed_variants() {
let routed = route_jsonrpc_message(JSONRPCMessage::Notification(JSONRPCNotification {
method: INITIALIZED_METHOD.to_string(),
params: Some(json!({})),
}))
.expect("initialized notification should route");
assert_eq!(
routed,
RoutedExecServerMessage::Inbound(ExecServerInboundMessage::Notification(
ExecServerClientNotification::Initialized,
))
);
}
#[test]
fn serializes_typed_notifications_back_to_jsonrpc() {
let message = encode_outbound_message(ExecServerOutboundMessage::Notification(
ExecServerServerNotification::Exited(ExecExitedNotification {
process_id: "proc-1".to_string(),
exit_code: 0,
}),
))
.expect("notification should serialize");
assert_eq!(
message,
JSONRPCMessage::Notification(JSONRPCNotification {
method: EXEC_EXITED_METHOD.to_string(),
params: Some(json!({
"processId": "proc-1",
"exitCode": 0,
})),
})
);
}
#[test]
fn serializes_typed_responses_back_to_jsonrpc() {
let message = encode_outbound_message(ExecServerOutboundMessage::Response {
request_id: RequestId::Integer(3),
response: ExecServerResponseMessage::Exec(ExecResponse {
process_id: "proc-1".to_string(),
}),
})
.expect("response should serialize");
assert_eq!(
message,
JSONRPCMessage::Response(codex_app_server_protocol::JSONRPCResponse {
id: RequestId::Integer(3),
result: json!({
"processId": "proc-1",
}),
})
);
}
#[test]
fn routes_exec_requests_with_typed_params() {
let cwd = std::env::current_dir().expect("cwd");
let routed = route_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest {
id: RequestId::Integer(4),
method: EXEC_METHOD.to_string(),
params: Some(json!({
"processId": "proc-1",
"argv": ["bash", "-lc", "true"],
"cwd": cwd,
"env": {},
"tty": true,
"arg0": null,
})),
trace: None,
}))
.expect("exec request should route");
let RoutedExecServerMessage::Inbound(ExecServerInboundMessage::Request(
ExecServerRequest::Exec { request_id, params },
)) = routed
else {
panic!("expected typed exec request");
};
assert_eq!(request_id, RequestId::Integer(4));
assert_eq!(
params,
ExecParams {
process_id: "proc-1".to_string(),
argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()],
cwd: std::env::current_dir().expect("cwd"),
env: std::collections::HashMap::new(),
tty: true,
arg0: None,
sandbox: None,
}
);
}
#[test]
fn routes_exec_requests_with_optional_sandbox_config() {
let cwd = std::env::current_dir().expect("cwd");
let routed = route_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest {
id: RequestId::Integer(4),
method: EXEC_METHOD.to_string(),
params: Some(json!({
"processId": "proc-1",
"argv": ["bash", "-lc", "true"],
"cwd": cwd,
"env": {},
"tty": true,
"arg0": null,
"sandbox": {
"mode": "none",
},
})),
trace: None,
}))
.expect("exec request with sandbox should route");
let RoutedExecServerMessage::Inbound(ExecServerInboundMessage::Request(
ExecServerRequest::Exec { request_id, params },
)) = routed
else {
panic!("expected typed exec request");
};
assert_eq!(request_id, RequestId::Integer(4));
assert_eq!(
params,
ExecParams {
process_id: "proc-1".to_string(),
argv: vec!["bash".to_string(), "-lc".to_string(), "true".to_string()],
cwd: std::env::current_dir().expect("cwd"),
env: std::collections::HashMap::new(),
tty: true,
arg0: None,
sandbox: Some(ExecSandboxConfig {
mode: ExecSandboxMode::None,
}),
}
);
}
#[test]
fn unknown_request_methods_return_immediate_invalid_request_errors() {
let routed = route_jsonrpc_message(JSONRPCMessage::Request(JSONRPCRequest {
id: RequestId::Integer(5),
method: "process/unknown".to_string(),
params: Some(json!({})),
trace: None,
}))
.expect("unknown request should still route");
assert_eq!(
routed,
RoutedExecServerMessage::ImmediateOutbound(ExecServerOutboundMessage::Error {
request_id: RequestId::Integer(5),
error: super::invalid_request("unknown method: process/unknown".to_string()),
})
);
}
#[test]
fn unexpected_client_notifications_are_rejected() {
let err = route_jsonrpc_message(JSONRPCMessage::Notification(JSONRPCNotification {
method: "process/output".to_string(),
params: Some(json!({})),
}))
.expect_err("unexpected client notification should fail");
assert_eq!(err, "unexpected notification method: process/output");
}
#[test]
fn unexpected_client_responses_are_rejected() {
let err = route_jsonrpc_message(JSONRPCMessage::Response(JSONRPCResponse {
id: RequestId::Integer(6),
result: json!({}),
}))
.expect_err("unexpected client response should fail");
assert_eq!(err, "unexpected client response for request id Integer(6)");
}
}

55
codex-rs/exec-server/src/server/transport.rs
View File

@@ -6,11 +6,11 @@ use tokio_tungstenite::accept_async;
use tracing::warn;
use crate::connection::JsonRpcConnection;
use crate::server::ExecServerConfig;
use crate::server::processor::run_connection;
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum ExecServerTransport {
Stdio,
WebSocket { bind_address: SocketAddr },
}
@@ -25,7 +25,7 @@ impl std::fmt::Display for ExecServerTransportParseError {
match self {
ExecServerTransportParseError::UnsupportedListenUrl(listen_url) => write!(
f,
"unsupported --listen URL `{listen_url}`; expected `stdio://` or `ws://IP:PORT`"
"unsupported --listen URL `{listen_url}`; expected `ws://IP:PORT`"
),
ExecServerTransportParseError::InvalidWebSocketListenUrl(listen_url) => write!(
f,
@@ -38,13 +38,9 @@ impl std::fmt::Display for ExecServerTransportParseError {
impl std::error::Error for ExecServerTransportParseError {}
impl ExecServerTransport {
pub const DEFAULT_LISTEN_URL: &str = "stdio://";
pub const DEFAULT_LISTEN_URL: &str = "ws://127.0.0.1:0";
pub fn from_listen_url(listen_url: &str) -> Result<Self, ExecServerTransportParseError> {
if listen_url == Self::DEFAULT_LISTEN_URL {
return Ok(Self::Stdio);
}
if let Some(socket_addr) = listen_url.strip_prefix("ws://") {
let bind_address = socket_addr.parse::<SocketAddr>().map_err(|_| {
ExecServerTransportParseError::InvalidWebSocketListenUrl(listen_url.to_string())
@@ -58,6 +54,16 @@ impl ExecServerTransport {
}
}
impl Default for ExecServerTransport {
fn default() -> Self {
Self::WebSocket {
bind_address: "127.0.0.1:0".parse().unwrap_or_else(|err| {
panic!("default exec-server websocket bind address should parse: {err}")
}),
}
}
}
impl FromStr for ExecServerTransport {
type Err = ExecServerTransportParseError;
@@ -68,25 +74,18 @@ impl FromStr for ExecServerTransport {
pub(crate) async fn run_transport(
transport: ExecServerTransport,
config: ExecServerConfig,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
match transport {
ExecServerTransport::Stdio => {
run_connection(JsonRpcConnection::from_stdio(
tokio::io::stdin(),
tokio::io::stdout(),
"exec-server stdio".to_string(),
))
.await;
Ok(())
}
ExecServerTransport::WebSocket { bind_address } => {
run_websocket_listener(bind_address).await
run_websocket_listener(bind_address, config).await
}
}
}
async fn run_websocket_listener(
bind_address: SocketAddr,
config: ExecServerConfig,
) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
let listener = TcpListener::bind(bind_address).await?;
let local_addr = listener.local_addr()?;
@@ -94,13 +93,17 @@ async fn run_websocket_listener(
loop {
let (stream, peer_addr) = listener.accept().await?;
let config = config.clone();
tokio::spawn(async move {
match accept_async(stream).await {
Ok(websocket) => {
run_connection(JsonRpcConnection::from_websocket(
websocket,
format!("exec-server websocket {peer_addr}"),
))
run_connection(
JsonRpcConnection::from_websocket(
websocket,
format!("exec-server websocket {peer_addr}"),
),
config,
)
.await;
}
Err(err) => {
@@ -124,14 +127,6 @@ mod tests {
use super::ExecServerTransport;
#[test]
fn exec_server_transport_parses_stdio_listen_url() {
let transport =
ExecServerTransport::from_listen_url(ExecServerTransport::DEFAULT_LISTEN_URL)
.expect("stdio listen URL should parse");
assert_eq!(transport, ExecServerTransport::Stdio);
}
#[test]
fn exec_server_transport_parses_websocket_listen_url() {
let transport = ExecServerTransport::from_listen_url("ws://127.0.0.1:1234")
@@ -160,7 +155,7 @@ mod tests {
.expect_err("unsupported scheme should fail");
assert_eq!(
err.to_string(),
"unsupported --listen URL `http://127.0.0.1:1234`; expected `stdio://` or `ws://IP:PORT`"
"unsupported --listen URL `http://127.0.0.1:1234`; expected `ws://IP:PORT`"
);
}
}