Merge remote-tracking branch 'origin/main' into pakrym/async-fs-trust-migration-pr

# Conflicts:
#	codex-rs/codex-mcp/src/mcp/mod_tests.rs
#	codex-rs/core/src/config/config_tests.rs
#	codex-rs/core/src/config/mod.rs
#	codex-rs/core/src/config_loader/mod.rs
#	codex-rs/core/src/guardian/review_session.rs
#	codex-rs/core/src/original_image_detail_tests.rs
#	codex-rs/core/src/realtime_context_tests.rs
#	codex-rs/core/src/tools/spec_tests.rs
#	codex-rs/exec-server/src/environment.rs
#	codex-rs/exec-server/src/file_system.rs
#	codex-rs/exec-server/src/lib.rs
#	codex-rs/models-manager/src/model_info_tests.rs
#	codex-rs/tui_app_server/src/onboarding/onboarding_screen.rs
#	codex-rs/tui_app_server/src/onboarding/trust_directory.rs
This commit is contained in:
pakrym-oai
2026-04-07 14:42:16 -07:00
2409 changed files with 97955 additions and 193715 deletions
+1
View File
@@ -15,6 +15,7 @@ path = "src/bin/codex-exec-server.rs"
workspace = true
[dependencies]
arc-swap = { workspace = true }
async-trait = { workspace = true }
base64 = { workspace = true }
clap = { workspace = true, features = ["derive"] }
+405 -18
View File
@@ -1,6 +1,8 @@
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;
use arc_swap::ArcSwap;
use codex_app_server_protocol::FsCopyParams;
use codex_app_server_protocol::FsCopyResponse;
use codex_app_server_protocol::FsCreateDirectoryParams;
@@ -17,22 +19,25 @@ use codex_app_server_protocol::FsWriteFileParams;
use codex_app_server_protocol::FsWriteFileResponse;
use codex_app_server_protocol::JSONRPCNotification;
use serde_json::Value;
use tokio::sync::broadcast;
use tokio::sync::Mutex;
use tokio::sync::watch;
use tokio::time::timeout;
use tokio_tungstenite::connect_async;
use tracing::debug;
use tracing::warn;
use crate::ProcessId;
use crate::client_api::ExecServerClientConnectOptions;
use crate::client_api::RemoteExecServerConnectArgs;
use crate::connection::JsonRpcConnection;
use crate::process::ExecServerEvent;
use crate::protocol::EXEC_CLOSED_METHOD;
use crate::protocol::EXEC_EXITED_METHOD;
use crate::protocol::EXEC_METHOD;
use crate::protocol::EXEC_OUTPUT_DELTA_METHOD;
use crate::protocol::EXEC_READ_METHOD;
use crate::protocol::EXEC_TERMINATE_METHOD;
use crate::protocol::EXEC_WRITE_METHOD;
use crate::protocol::ExecClosedNotification;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecParams;
@@ -90,9 +95,29 @@ impl RemoteExecServerConnectArgs {
}
}
pub(crate) struct SessionState {
wake_tx: watch::Sender<u64>,
failure: Mutex<Option<String>>,
}
#[derive(Clone)]
pub(crate) struct Session {
client: ExecServerClient,
process_id: ProcessId,
state: Arc<SessionState>,
}
struct Inner {
client: RpcClient,
events_tx: broadcast::Sender<ExecServerEvent>,
// The remote transport delivers one shared notification stream for every
// process on the connection. Keep a local process_id -> session registry so
// we can turn those connection-global notifications into process wakeups
// without making notifications the source of truth for output delivery.
sessions: ArcSwap<HashMap<ProcessId, Arc<SessionState>>>,
// ArcSwap makes reads cheap on the hot notification path, but writes still
// need serialization so concurrent register/remove operations do not
// overwrite each other's copy-on-write updates.
sessions_write_lock: Mutex<()>,
reader_task: tokio::task::JoinHandle<()>,
}
@@ -158,10 +183,6 @@ impl ExecServerClient {
.await
}
pub fn event_receiver(&self) -> broadcast::Receiver<ExecServerEvent> {
self.inner.events_tx.subscribe()
}
pub async fn initialize(
&self,
options: ExecServerClientConnectOptions,
@@ -204,7 +225,7 @@ impl ExecServerClient {
pub async fn write(
&self,
process_id: &str,
process_id: &ProcessId,
chunk: Vec<u8>,
) -> Result<WriteResponse, ExecServerError> {
self.inner
@@ -212,7 +233,7 @@ impl ExecServerClient {
.call(
EXEC_WRITE_METHOD,
&WriteParams {
process_id: process_id.to_string(),
process_id: process_id.clone(),
chunk: chunk.into(),
},
)
@@ -220,13 +241,16 @@ impl ExecServerClient {
.map_err(Into::into)
}
pub async fn terminate(&self, process_id: &str) -> Result<TerminateResponse, ExecServerError> {
pub async fn terminate(
&self,
process_id: &ProcessId,
) -> Result<TerminateResponse, ExecServerError> {
self.inner
.client
.call(
EXEC_TERMINATE_METHOD,
&TerminateParams {
process_id: process_id.to_string(),
process_id: process_id.clone(),
},
)
.await
@@ -307,6 +331,25 @@ impl ExecServerClient {
.map_err(Into::into)
}
pub(crate) async fn register_session(
&self,
process_id: &ProcessId,
) -> Result<Session, ExecServerError> {
let state = Arc::new(SessionState::new());
self.inner
.insert_session(process_id, Arc::clone(&state))
.await?;
Ok(Session {
client: self.clone(),
process_id: process_id.clone(),
state,
})
}
pub(crate) async fn unregister_session(&self, process_id: &ProcessId) {
self.inner.remove_session(process_id).await;
}
async fn connect(
connection: JsonRpcConnection,
options: ExecServerClientConnectOptions,
@@ -322,13 +365,18 @@ impl ExecServerClient {
&& let Err(err) =
handle_server_notification(&inner, notification).await
{
warn!("exec-server client closing after protocol error: {err}");
fail_all_sessions(
&inner,
format!("exec-server notification handling failed: {err}"),
)
.await;
return;
}
}
RpcClientEvent::Disconnected { reason } => {
if let Some(reason) = reason {
warn!("exec-server client transport disconnected: {reason}");
if let Some(inner) = weak.upgrade() {
fail_all_sessions(&inner, disconnected_message(reason.as_deref()))
.await;
}
return;
}
@@ -338,7 +386,8 @@ impl ExecServerClient {
Inner {
client: rpc_client,
events_tx: broadcast::channel(256).0,
sessions: ArcSwap::from_pointee(HashMap::new()),
sessions_write_lock: Mutex::new(()),
reader_task,
}
});
@@ -370,6 +419,177 @@ impl From<RpcCallError> for ExecServerError {
}
}
impl SessionState {
fn new() -> Self {
let (wake_tx, _wake_rx) = watch::channel(0);
Self {
wake_tx,
failure: Mutex::new(None),
}
}
pub(crate) fn subscribe(&self) -> watch::Receiver<u64> {
self.wake_tx.subscribe()
}
fn note_change(&self, seq: u64) {
let next = (*self.wake_tx.borrow()).max(seq);
let _ = self.wake_tx.send(next);
}
async fn set_failure(&self, message: String) {
let mut failure = self.failure.lock().await;
if failure.is_none() {
*failure = Some(message);
}
drop(failure);
let next = (*self.wake_tx.borrow()).saturating_add(1);
let _ = self.wake_tx.send(next);
}
async fn failed_response(&self) -> Option<ReadResponse> {
self.failure
.lock()
.await
.clone()
.map(|message| self.synthesized_failure(message))
}
fn synthesized_failure(&self, message: String) -> ReadResponse {
let next_seq = (*self.wake_tx.borrow()).saturating_add(1);
ReadResponse {
chunks: Vec::new(),
next_seq,
exited: true,
exit_code: None,
closed: true,
failure: Some(message),
}
}
}
impl Session {
pub(crate) fn process_id(&self) -> &ProcessId {
&self.process_id
}
pub(crate) fn subscribe_wake(&self) -> watch::Receiver<u64> {
self.state.subscribe()
}
pub(crate) async fn read(
&self,
after_seq: Option<u64>,
max_bytes: Option<usize>,
wait_ms: Option<u64>,
) -> Result<ReadResponse, ExecServerError> {
if let Some(response) = self.state.failed_response().await {
return Ok(response);
}
match self
.client
.read(ReadParams {
process_id: self.process_id.clone(),
after_seq,
max_bytes,
wait_ms,
})
.await
{
Ok(response) => Ok(response),
Err(err) if is_transport_closed_error(&err) => {
let message = disconnected_message(/*reason*/ None);
self.state.set_failure(message.clone()).await;
Ok(self.state.synthesized_failure(message))
}
Err(err) => Err(err),
}
}
pub(crate) async fn write(&self, chunk: Vec<u8>) -> Result<WriteResponse, ExecServerError> {
self.client.write(&self.process_id, chunk).await
}
pub(crate) async fn terminate(&self) -> Result<(), ExecServerError> {
self.client.terminate(&self.process_id).await?;
Ok(())
}
pub(crate) async fn unregister(&self) {
self.client.unregister_session(&self.process_id).await;
}
}
impl Inner {
fn get_session(&self, process_id: &ProcessId) -> Option<Arc<SessionState>> {
self.sessions.load().get(process_id).cloned()
}
async fn insert_session(
&self,
process_id: &ProcessId,
session: Arc<SessionState>,
) -> Result<(), ExecServerError> {
let _sessions_write_guard = self.sessions_write_lock.lock().await;
let sessions = self.sessions.load();
if sessions.contains_key(process_id) {
return Err(ExecServerError::Protocol(format!(
"session already registered for process {process_id}"
)));
}
let mut next_sessions = sessions.as_ref().clone();
next_sessions.insert(process_id.clone(), session);
self.sessions.store(Arc::new(next_sessions));
Ok(())
}
async fn remove_session(&self, process_id: &ProcessId) -> Option<Arc<SessionState>> {
let _sessions_write_guard = self.sessions_write_lock.lock().await;
let sessions = self.sessions.load();
let session = sessions.get(process_id).cloned();
session.as_ref()?;
let mut next_sessions = sessions.as_ref().clone();
next_sessions.remove(process_id);
self.sessions.store(Arc::new(next_sessions));
session
}
async fn take_all_sessions(&self) -> HashMap<ProcessId, Arc<SessionState>> {
let _sessions_write_guard = self.sessions_write_lock.lock().await;
let sessions = self.sessions.load();
let drained_sessions = sessions.as_ref().clone();
self.sessions.store(Arc::new(HashMap::new()));
drained_sessions
}
}
fn disconnected_message(reason: Option<&str>) -> String {
match reason {
Some(reason) => format!("exec-server transport disconnected: {reason}"),
None => "exec-server transport disconnected".to_string(),
}
}
fn is_transport_closed_error(error: &ExecServerError) -> bool {
matches!(error, ExecServerError::Closed)
|| matches!(
error,
ExecServerError::Server {
code: -32000,
message,
} if message == "JSON-RPC transport closed"
)
}
async fn fail_all_sessions(inner: &Arc<Inner>, message: String) {
let sessions = inner.take_all_sessions().await;
for (_, session) in sessions {
session.set_failure(message.clone()).await;
}
}
async fn handle_server_notification(
inner: &Arc<Inner>,
notification: JSONRPCNotification,
@@ -378,12 +598,26 @@ async fn handle_server_notification(
EXEC_OUTPUT_DELTA_METHOD => {
let params: ExecOutputDeltaNotification =
serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
let _ = inner.events_tx.send(ExecServerEvent::OutputDelta(params));
if let Some(session) = inner.get_session(&params.process_id) {
session.note_change(params.seq);
}
}
EXEC_EXITED_METHOD => {
let params: ExecExitedNotification =
serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
let _ = inner.events_tx.send(ExecServerEvent::Exited(params));
if let Some(session) = inner.get_session(&params.process_id) {
session.note_change(params.seq);
}
}
EXEC_CLOSED_METHOD => {
let params: ExecClosedNotification =
serde_json::from_value(notification.params.unwrap_or(Value::Null))?;
// Closed is the terminal lifecycle event for this process, so drop
// the routing entry before forwarding it.
let session = inner.remove_session(&params.process_id).await;
if let Some(session) = session {
session.note_change(params.seq);
}
}
other => {
debug!("ignoring unknown exec-server notification: {other}");
@@ -391,3 +625,156 @@ async fn handle_server_notification(
}
Ok(())
}
#[cfg(test)]
mod tests {
use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCNotification;
use codex_app_server_protocol::JSONRPCResponse;
use pretty_assertions::assert_eq;
use tokio::io::AsyncBufReadExt;
use tokio::io::AsyncWrite;
use tokio::io::AsyncWriteExt;
use tokio::io::BufReader;
use tokio::io::duplex;
use tokio::sync::mpsc;
use tokio::time::Duration;
use tokio::time::timeout;
use super::ExecServerClient;
use super::ExecServerClientConnectOptions;
use crate::ProcessId;
use crate::connection::JsonRpcConnection;
use crate::protocol::EXEC_EXITED_METHOD;
use crate::protocol::EXEC_OUTPUT_DELTA_METHOD;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecOutputStream;
use crate::protocol::INITIALIZE_METHOD;
use crate::protocol::INITIALIZED_METHOD;
use crate::protocol::InitializeResponse;
async fn read_jsonrpc_line<R>(lines: &mut tokio::io::Lines<BufReader<R>>) -> JSONRPCMessage
where
R: tokio::io::AsyncRead + Unpin,
{
let line = timeout(Duration::from_secs(1), lines.next_line())
.await
.expect("json-rpc read should not time out")
.expect("json-rpc read should succeed")
.expect("json-rpc connection should stay open");
serde_json::from_str(&line).expect("json-rpc line should parse")
}
async fn write_jsonrpc_line<W>(writer: &mut W, message: JSONRPCMessage)
where
W: AsyncWrite + Unpin,
{
let encoded = serde_json::to_string(&message).expect("json-rpc message should serialize");
writer
.write_all(format!("{encoded}\n").as_bytes())
.await
.expect("json-rpc line should write");
}
#[tokio::test]
async fn wake_notifications_do_not_block_other_sessions() {
let (client_stdin, server_reader) = duplex(1 << 20);
let (mut server_writer, client_stdout) = duplex(1 << 20);
let (notifications_tx, mut notifications_rx) = mpsc::channel(16);
let server = tokio::spawn(async move {
let mut lines = BufReader::new(server_reader).lines();
let initialize = read_jsonrpc_line(&mut lines).await;
let request = match initialize {
JSONRPCMessage::Request(request) if request.method == INITIALIZE_METHOD => request,
other => panic!("expected initialize request, got {other:?}"),
};
write_jsonrpc_line(
&mut server_writer,
JSONRPCMessage::Response(JSONRPCResponse {
id: request.id,
result: serde_json::to_value(InitializeResponse {})
.expect("initialize response should serialize"),
}),
)
.await;
let initialized = read_jsonrpc_line(&mut lines).await;
match initialized {
JSONRPCMessage::Notification(notification)
if notification.method == INITIALIZED_METHOD => {}
other => panic!("expected initialized notification, got {other:?}"),
}
while let Some(message) = notifications_rx.recv().await {
write_jsonrpc_line(&mut server_writer, message).await;
}
});
let client = ExecServerClient::connect(
JsonRpcConnection::from_stdio(
client_stdout,
client_stdin,
"test-exec-server-client".to_string(),
),
ExecServerClientConnectOptions::default(),
)
.await
.expect("client should connect");
let noisy_process_id = ProcessId::from("noisy");
let quiet_process_id = ProcessId::from("quiet");
let _noisy_session = client
.register_session(&noisy_process_id)
.await
.expect("noisy session should register");
let quiet_session = client
.register_session(&quiet_process_id)
.await
.expect("quiet session should register");
let mut quiet_wake_rx = quiet_session.subscribe_wake();
for seq in 0..=4096 {
notifications_tx
.send(JSONRPCMessage::Notification(JSONRPCNotification {
method: EXEC_OUTPUT_DELTA_METHOD.to_string(),
params: Some(
serde_json::to_value(ExecOutputDeltaNotification {
process_id: noisy_process_id.clone(),
seq,
stream: ExecOutputStream::Stdout,
chunk: b"x".to_vec().into(),
})
.expect("output notification should serialize"),
),
}))
.await
.expect("output notification should queue");
}
notifications_tx
.send(JSONRPCMessage::Notification(JSONRPCNotification {
method: EXEC_EXITED_METHOD.to_string(),
params: Some(
serde_json::to_value(ExecExitedNotification {
process_id: quiet_process_id,
seq: 1,
exit_code: 17,
})
.expect("exit notification should serialize"),
),
}))
.await
.expect("exit notification should queue");
timeout(Duration::from_secs(1), quiet_wake_rx.changed())
.await
.expect("quiet session should receive wake before timeout")
.expect("quiet wake channel should stay open");
assert_eq!(*quiet_wake_rx.borrow(), 1);
drop(notifications_tx);
drop(client);
server.await.expect("server task should finish");
}
}
+186 -53
View File
@@ -1,24 +1,101 @@
use std::sync::Arc;
use tokio::sync::OnceCell;
use crate::ExecServerClient;
use crate::ExecServerError;
use crate::LOCAL_FS;
use crate::RemoteExecServerConnectArgs;
use crate::file_system::ExecutorFileSystem;
use crate::local_process::LocalProcess;
use crate::process::ExecProcess;
use crate::process::ExecBackend;
use crate::remote_file_system::RemoteFileSystem;
use crate::remote_process::RemoteProcess;
pub trait ExecutorEnvironment: Send + Sync {
fn get_executor(&self) -> Arc<dyn ExecProcess>;
pub const CODEX_EXEC_SERVER_URL_ENV_VAR: &str = "CODEX_EXEC_SERVER_URL";
/// Lazily creates and caches the active environment for a session.
///
/// The manager keeps the session's environment selection stable so subagents
/// and follow-up turns preserve an explicit disabled state.
#[derive(Debug)]
pub struct EnvironmentManager {
exec_server_url: Option<String>,
disabled: bool,
current_environment: OnceCell<Option<Arc<Environment>>>,
}
impl Default for EnvironmentManager {
fn default() -> Self {
Self::new(/*exec_server_url*/ None)
}
}
impl EnvironmentManager {
/// Builds a manager from the raw `CODEX_EXEC_SERVER_URL` value.
pub fn new(exec_server_url: Option<String>) -> Self {
let (exec_server_url, disabled) = normalize_exec_server_url(exec_server_url);
Self {
exec_server_url,
disabled,
current_environment: OnceCell::new(),
}
}
/// Builds a manager from process environment variables.
pub fn from_env() -> Self {
Self::new(std::env::var(CODEX_EXEC_SERVER_URL_ENV_VAR).ok())
}
/// Builds a manager from the currently selected environment, or from the
/// disabled mode when no environment is available.
pub fn from_environment(environment: Option<&Environment>) -> Self {
match environment {
Some(environment) => Self {
exec_server_url: environment.exec_server_url().map(str::to_owned),
disabled: false,
current_environment: OnceCell::new(),
},
None => Self {
exec_server_url: None,
disabled: true,
current_environment: OnceCell::new(),
},
}
}
/// Returns the remote exec-server URL when one is configured.
pub fn exec_server_url(&self) -> Option<&str> {
self.exec_server_url.as_deref()
}
/// Returns the cached environment, creating it on first access.
pub async fn current(&self) -> Result<Option<Arc<Environment>>, ExecServerError> {
self.current_environment
.get_or_try_init(|| async {
if self.disabled {
Ok(None)
} else {
Ok(Some(Arc::new(
Environment::create(self.exec_server_url.clone()).await?,
)))
}
})
.await
.map(Option::as_ref)
.map(std::option::Option::<&Arc<Environment>>::cloned)
}
}
/// Concrete execution/filesystem environment selected for a session.
///
/// This bundles the selected backend together with the corresponding remote
/// client, if any.
#[derive(Clone)]
pub struct Environment {
experimental_exec_server_url: Option<String>,
exec_server_url: Option<String>,
remote_exec_server_client: Option<ExecServerClient>,
executor: Arc<dyn ExecProcess>,
exec_backend: Arc<dyn ExecBackend>,
}
impl Default for Environment {
@@ -32,9 +109,9 @@ impl Default for Environment {
}
Self {
experimental_exec_server_url: None,
exec_server_url: None,
remote_exec_server_client: None,
executor: Arc::new(local_process),
exec_backend: Arc::new(local_process),
}
}
}
@@ -42,22 +119,25 @@ impl Default for Environment {
impl std::fmt::Debug for Environment {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("Environment")
.field(
"experimental_exec_server_url",
&self.experimental_exec_server_url,
)
.field("exec_server_url", &self.exec_server_url)
.finish_non_exhaustive()
}
}
impl Environment {
pub async fn create(
experimental_exec_server_url: Option<String>,
) -> Result<Self, ExecServerError> {
let remote_exec_server_client = if let Some(url) = &experimental_exec_server_url {
/// Builds an environment from the raw `CODEX_EXEC_SERVER_URL` value.
pub async fn create(exec_server_url: Option<String>) -> Result<Self, ExecServerError> {
let (exec_server_url, disabled) = normalize_exec_server_url(exec_server_url);
if disabled {
return Err(ExecServerError::Protocol(
"disabled mode does not create an Environment".to_string(),
));
}
let remote_exec_server_client = if let Some(exec_server_url) = &exec_server_url {
Some(
ExecServerClient::connect_websocket(RemoteExecServerConnectArgs {
websocket_url: url.clone(),
websocket_url: exec_server_url.clone(),
client_name: "codex-environment".to_string(),
connect_timeout: std::time::Duration::from_secs(5),
initialize_timeout: std::time::Duration::from_secs(5),
@@ -68,71 +148,129 @@ impl Environment {
None
};
let executor: Arc<dyn ExecProcess> = if let Some(client) = remote_exec_server_client.clone()
{
Arc::new(RemoteProcess::new(client))
} else {
let local_process = LocalProcess::default();
local_process
.initialize()
.map_err(|err| ExecServerError::Protocol(err.message))?;
local_process
.initialized()
.map_err(ExecServerError::Protocol)?;
Arc::new(local_process)
let exec_backend: Arc<dyn ExecBackend> = match remote_exec_server_client.clone() {
Some(client) => Arc::new(RemoteProcess::new(client)),
None if exec_server_url.is_some() => {
return Err(ExecServerError::Protocol(
"remote mode should have an exec-server client".to_string(),
));
}
None => {
let local_process = LocalProcess::default();
local_process
.initialize()
.map_err(|err| ExecServerError::Protocol(err.message))?;
local_process
.initialized()
.map_err(ExecServerError::Protocol)?;
Arc::new(local_process)
}
};
Ok(Self {
experimental_exec_server_url,
exec_server_url,
remote_exec_server_client,
executor,
exec_backend,
})
}
pub fn experimental_exec_server_url(&self) -> Option<&str> {
self.experimental_exec_server_url.as_deref()
pub fn is_remote(&self) -> bool {
self.exec_server_url.is_some()
}
pub fn get_executor(&self) -> Arc<dyn ExecProcess> {
Arc::clone(&self.executor)
/// Returns the remote exec-server URL when this environment is remote.
pub fn exec_server_url(&self) -> Option<&str> {
self.exec_server_url.as_deref()
}
pub fn get_exec_backend(&self) -> Arc<dyn ExecBackend> {
Arc::clone(&self.exec_backend)
}
pub fn get_filesystem(&self) -> Arc<dyn ExecutorFileSystem> {
if let Some(client) = self.remote_exec_server_client.clone() {
Arc::new(RemoteFileSystem::new(client))
} else {
Arc::new(LOCAL_FS)
match self.remote_exec_server_client.clone() {
Some(client) => Arc::new(RemoteFileSystem::new(client)),
None => Arc::clone(&LOCAL_FS),
}
}
}
impl ExecutorEnvironment for Environment {
fn get_executor(&self) -> Arc<dyn ExecProcess> {
Arc::clone(&self.executor)
fn normalize_exec_server_url(exec_server_url: Option<String>) -> (Option<String>, bool) {
match exec_server_url.as_deref().map(str::trim) {
None | Some("") => (None, false),
Some(url) if url.eq_ignore_ascii_case("none") => (None, true),
Some(url) => (Some(url.to_string()), false),
}
}
#[cfg(test)]
mod tests {
use std::sync::Arc;
use super::Environment;
use super::EnvironmentManager;
use crate::ProcessId;
use pretty_assertions::assert_eq;
#[tokio::test]
async fn create_without_remote_exec_server_url_does_not_connect() {
let environment = Environment::create(None).await.expect("create environment");
async fn create_local_environment_does_not_connect() {
let environment = Environment::create(/*exec_server_url*/ None)
.await
.expect("create environment");
assert_eq!(environment.experimental_exec_server_url(), None);
assert_eq!(environment.exec_server_url(), None);
assert!(environment.remote_exec_server_client.is_none());
}
#[test]
fn environment_manager_normalizes_empty_url() {
let manager = EnvironmentManager::new(Some(String::new()));
assert!(!manager.disabled);
assert_eq!(manager.exec_server_url(), None);
}
#[test]
fn environment_manager_treats_none_value_as_disabled() {
let manager = EnvironmentManager::new(Some("none".to_string()));
assert!(manager.disabled);
assert_eq!(manager.exec_server_url(), None);
}
#[tokio::test]
async fn environment_manager_current_caches_environment() {
let manager = EnvironmentManager::new(/*exec_server_url*/ None);
let first = manager.current().await.expect("get current environment");
let second = manager.current().await.expect("get current environment");
let first = first.expect("local environment");
let second = second.expect("local environment");
assert!(Arc::ptr_eq(&first, &second));
}
#[tokio::test]
async fn disabled_environment_manager_has_no_current_environment() {
let manager = EnvironmentManager::new(Some("none".to_string()));
assert!(
manager
.current()
.await
.expect("get current environment")
.is_none()
);
}
#[tokio::test]
async fn default_environment_has_ready_local_executor() {
let environment = Environment::default();
let response = environment
.get_executor()
.get_exec_backend()
.start(crate::ExecParams {
process_id: "default-env-proc".to_string(),
process_id: ProcessId::from("default-env-proc"),
argv: vec!["true".to_string()],
cwd: std::env::current_dir().expect("read current dir"),
env: Default::default(),
@@ -142,11 +280,6 @@ mod tests {
.await
.expect("start process");
assert_eq!(
response,
crate::ExecResponse {
process_id: "default-env-proc".to_string(),
}
);
assert_eq!(response.process.process_id().as_str(), "default-env-proc");
}
}
+2 -1
View File
@@ -39,7 +39,8 @@ pub type FileSystemResult<T> = io::Result<T>;
pub trait ExecutorFileSystem: Send + Sync {
async fn read_file(&self, path: &AbsolutePathBuf) -> FileSystemResult<Vec<u8>>;
async fn read_to_string(&self, path: &AbsolutePathBuf) -> FileSystemResult<String> {
/// Reads a file and decodes it as UTF-8 text.
async fn read_file_text(&self, path: &AbsolutePathBuf) -> FileSystemResult<String> {
let bytes = self.read_file(path).await?;
String::from_utf8(bytes).map_err(|err| io::Error::new(io::ErrorKind::InvalidData, err))
}
+9 -5
View File
@@ -6,6 +6,7 @@ mod file_system;
mod local_file_system;
mod local_process;
mod process;
mod process_id;
mod protocol;
mod remote_file_system;
mod remote_process;
@@ -30,8 +31,9 @@ pub use codex_app_server_protocol::FsRemoveParams;
pub use codex_app_server_protocol::FsRemoveResponse;
pub use codex_app_server_protocol::FsWriteFileParams;
pub use codex_app_server_protocol::FsWriteFileResponse;
pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR;
pub use environment::Environment;
pub use environment::ExecutorEnvironment;
pub use environment::EnvironmentManager;
pub use file_system::CopyOptions;
pub use file_system::CreateDirectoryOptions;
pub use file_system::ExecutorFileSystem;
@@ -39,9 +41,12 @@ pub use file_system::FileMetadata;
pub use file_system::FileSystemResult;
pub use file_system::ReadDirectoryEntry;
pub use file_system::RemoveOptions;
pub use local_file_system::LocalFileSystem;
pub use local_file_system::LOCAL_FS;
pub use process::ExecBackend;
pub use process::ExecProcess;
pub use process::ExecServerEvent;
pub use process::StartedExecProcess;
pub use process_id::ProcessId;
pub use protocol::ExecClosedNotification;
pub use protocol::ExecExitedNotification;
pub use protocol::ExecOutputDeltaNotification;
pub use protocol::ExecOutputStream;
@@ -55,9 +60,8 @@ pub use protocol::TerminateParams;
pub use protocol::TerminateResponse;
pub use protocol::WriteParams;
pub use protocol::WriteResponse;
pub use protocol::WriteStatus;
pub use server::DEFAULT_LISTEN_URL;
pub use server::ExecServerListenUrlParseError;
pub use server::run_main;
pub use server::run_main_with_listen_url;
pub const LOCAL_FS: LocalFileSystem = LocalFileSystem;
@@ -3,6 +3,8 @@ use codex_utils_absolute_path::AbsolutePathBuf;
use std::path::Component;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::sync::LazyLock;
use std::time::SystemTime;
use std::time::UNIX_EPOCH;
use tokio::io;
@@ -17,6 +19,9 @@ use crate::RemoveOptions;
const MAX_READ_FILE_BYTES: u64 = 512 * 1024 * 1024;
pub static LOCAL_FS: LazyLock<Arc<dyn ExecutorFileSystem>> =
LazyLock::new(|| -> Arc<dyn ExecutorFileSystem> { Arc::new(LocalFileSystem) });
#[derive(Clone, Default)]
pub struct LocalFileSystem;
+198 -59
View File
@@ -11,13 +11,16 @@ use codex_utils_pty::ExecCommandSession;
use codex_utils_pty::TerminalSize;
use tokio::sync::Mutex;
use tokio::sync::Notify;
use tokio::sync::broadcast;
use tokio::sync::mpsc;
use tracing::warn;
use tokio::sync::watch;
use crate::ExecBackend;
use crate::ExecProcess;
use crate::ExecServerError;
use crate::ExecServerEvent;
use crate::ProcessId;
use crate::StartedExecProcess;
use crate::protocol::EXEC_CLOSED_METHOD;
use crate::protocol::ExecClosedNotification;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::protocol::ExecOutputStream;
@@ -31,6 +34,7 @@ use crate::protocol::TerminateParams;
use crate::protocol::TerminateResponse;
use crate::protocol::WriteParams;
use crate::protocol::WriteResponse;
use crate::protocol::WriteStatus;
use crate::rpc::RpcNotificationSender;
use crate::rpc::RpcServerOutboundMessage;
use crate::rpc::internal_error;
@@ -38,7 +42,6 @@ use crate::rpc::invalid_params;
use crate::rpc::invalid_request;
const RETAINED_OUTPUT_BYTES_PER_PROCESS: usize = 1024 * 1024;
const EVENT_CHANNEL_CAPACITY: usize = 256;
const NOTIFICATION_CHANNEL_CAPACITY: usize = 256;
#[cfg(test)]
const EXITED_PROCESS_RETENTION: Duration = Duration::from_millis(25);
@@ -59,7 +62,10 @@ struct RunningProcess {
retained_bytes: usize,
next_seq: u64,
exit_code: Option<i32>,
wake_tx: watch::Sender<u64>,
output_notify: Arc<Notify>,
open_streams: usize,
closed: bool,
}
enum ProcessEntry {
@@ -69,8 +75,7 @@ enum ProcessEntry {
struct Inner {
notifications: RpcNotificationSender,
events_tx: broadcast::Sender<ExecServerEvent>,
processes: Mutex<HashMap<String, ProcessEntry>>,
processes: Mutex<HashMap<ProcessId, ProcessEntry>>,
initialize_requested: AtomicBool,
initialized: AtomicBool,
}
@@ -80,6 +85,12 @@ pub(crate) struct LocalProcess {
inner: Arc<Inner>,
}
struct LocalExecProcess {
process_id: ProcessId,
backend: LocalProcess,
wake_tx: watch::Sender<u64>,
}
impl Default for LocalProcess {
fn default() -> Self {
let (outgoing_tx, mut outgoing_rx) =
@@ -94,7 +105,6 @@ impl LocalProcess {
Self {
inner: Arc::new(Inner {
notifications,
events_tx: broadcast::channel(EVENT_CHANNEL_CAPACITY).0,
processes: Mutex::new(HashMap::new()),
initialize_requested: AtomicBool::new(false),
initialized: AtomicBool::new(false),
@@ -152,10 +162,12 @@ impl LocalProcess {
Ok(())
}
pub(crate) async fn exec(&self, params: ExecParams) -> Result<ExecResponse, JSONRPCErrorError> {
async fn start_process(
&self,
params: ExecParams,
) -> Result<(ExecResponse, watch::Sender<u64>), JSONRPCErrorError> {
self.require_initialized_for("exec")?;
let process_id = params.process_id.clone();
let (program, args) = params
.argv
.split_first()
@@ -203,6 +215,7 @@ impl LocalProcess {
};
let output_notify = Arc::new(Notify::new());
let (wake_tx, _wake_rx) = watch::channel(0);
{
let mut process_map = self.inner.processes.lock().await;
process_map.insert(
@@ -214,7 +227,10 @@ impl LocalProcess {
retained_bytes: 0,
next_seq: 1,
exit_code: None,
wake_tx: wake_tx.clone(),
output_notify: Arc::clone(&output_notify),
open_streams: 2,
closed: false,
})),
);
}
@@ -248,7 +264,13 @@ impl LocalProcess {
output_notify,
));
Ok(ExecResponse { process_id })
Ok((ExecResponse { process_id }, wake_tx))
}
pub(crate) async fn exec(&self, params: ExecParams) -> Result<ExecResponse, JSONRPCErrorError> {
self.start_process(params)
.await
.map(|(response, _)| response)
}
pub(crate) async fn exec_read(
@@ -256,6 +278,7 @@ impl LocalProcess {
params: ReadParams,
) -> Result<ReadResponse, JSONRPCErrorError> {
self.require_initialized_for("exec")?;
let _process_id = params.process_id.clone();
let after_seq = params.after_seq.unwrap_or(0);
let max_bytes = params.max_bytes.unwrap_or(usize::MAX);
let wait = Duration::from_millis(params.wait_ms.unwrap_or(0));
@@ -300,6 +323,8 @@ impl LocalProcess {
next_seq,
exited: process.exit_code.is_some(),
exit_code: process.exit_code,
closed: process.closed,
failure: None,
},
Arc::clone(&process.output_notify),
)
@@ -309,6 +334,11 @@ impl LocalProcess {
|| response.exited
|| tokio::time::Instant::now() >= deadline
{
let _total_bytes: usize = response
.chunks
.iter()
.map(|chunk| chunk.chunk.0.len())
.sum();
return Ok(response);
}
@@ -325,22 +355,24 @@ impl LocalProcess {
params: WriteParams,
) -> Result<WriteResponse, JSONRPCErrorError> {
self.require_initialized_for("exec")?;
let _process_id = params.process_id.clone();
let _input_bytes = params.chunk.0.len();
let writer_tx = {
let process_map = self.inner.processes.lock().await;
let process = process_map.get(&params.process_id).ok_or_else(|| {
invalid_request(format!("unknown process id {}", params.process_id))
})?;
let Some(process) = process_map.get(&params.process_id) else {
return Ok(WriteResponse {
status: WriteStatus::UnknownProcess,
});
};
let ProcessEntry::Running(process) = process else {
return Err(invalid_request(format!(
"process id {} is starting",
params.process_id
)));
return Ok(WriteResponse {
status: WriteStatus::Starting,
});
};
if !process.tty {
return Err(invalid_request(format!(
"stdin is closed for process {}",
params.process_id
)));
return Ok(WriteResponse {
status: WriteStatus::StdinClosed,
});
}
process.session.writer_sender()
};
@@ -350,7 +382,9 @@ impl LocalProcess {
.await
.map_err(|_| internal_error("failed to write to process stdin".to_string()))?;
Ok(WriteResponse { accepted: true })
Ok(WriteResponse {
status: WriteStatus::Accepted,
})
}
pub(crate) async fn terminate_process(
@@ -358,6 +392,7 @@ impl LocalProcess {
params: TerminateParams,
) -> Result<TerminateResponse, JSONRPCErrorError> {
self.require_initialized_for("exec")?;
let _process_id = params.process_id.clone();
let running = {
let process_map = self.inner.processes.lock().await;
match process_map.get(&params.process_id) {
@@ -377,38 +412,90 @@ impl LocalProcess {
}
#[async_trait]
impl ExecProcess for LocalProcess {
async fn start(&self, params: ExecParams) -> Result<ExecResponse, ExecServerError> {
self.exec(params).await.map_err(map_handler_error)
impl ExecBackend for LocalProcess {
async fn start(&self, params: ExecParams) -> Result<StartedExecProcess, ExecServerError> {
let (response, wake_tx) = self
.start_process(params)
.await
.map_err(map_handler_error)?;
Ok(StartedExecProcess {
process: Arc::new(LocalExecProcess {
process_id: response.process_id,
backend: self.clone(),
wake_tx,
}),
})
}
}
#[async_trait]
impl ExecProcess for LocalExecProcess {
fn process_id(&self) -> &ProcessId {
&self.process_id
}
async fn read(&self, params: ReadParams) -> Result<ReadResponse, ExecServerError> {
self.exec_read(params).await.map_err(map_handler_error)
fn subscribe_wake(&self) -> watch::Receiver<u64> {
self.wake_tx.subscribe()
}
async fn read(
&self,
after_seq: Option<u64>,
max_bytes: Option<usize>,
wait_ms: Option<u64>,
) -> Result<ReadResponse, ExecServerError> {
self.backend
.read(&self.process_id, after_seq, max_bytes, wait_ms)
.await
}
async fn write(&self, chunk: Vec<u8>) -> Result<WriteResponse, ExecServerError> {
self.backend.write(&self.process_id, chunk).await
}
async fn terminate(&self) -> Result<(), ExecServerError> {
self.backend.terminate(&self.process_id).await
}
}
impl LocalProcess {
async fn read(
&self,
process_id: &ProcessId,
after_seq: Option<u64>,
max_bytes: Option<usize>,
wait_ms: Option<u64>,
) -> Result<ReadResponse, ExecServerError> {
self.exec_read(ReadParams {
process_id: process_id.clone(),
after_seq,
max_bytes,
wait_ms,
})
.await
.map_err(map_handler_error)
}
async fn write(
&self,
process_id: &str,
process_id: &ProcessId,
chunk: Vec<u8>,
) -> Result<WriteResponse, ExecServerError> {
self.exec_write(WriteParams {
process_id: process_id.to_string(),
process_id: process_id.clone(),
chunk: chunk.into(),
})
.await
.map_err(map_handler_error)
}
async fn terminate(&self, process_id: &str) -> Result<TerminateResponse, ExecServerError> {
async fn terminate(&self, process_id: &ProcessId) -> Result<(), ExecServerError> {
self.terminate_process(TerminateParams {
process_id: process_id.to_string(),
process_id: process_id.clone(),
})
.await
.map_err(map_handler_error)
}
fn subscribe_events(&self) -> broadcast::Receiver<ExecServerEvent> {
self.inner.events_tx.subscribe()
.map_err(map_handler_error)?;
Ok(())
}
}
@@ -420,13 +507,14 @@ fn map_handler_error(error: JSONRPCErrorError) -> ExecServerError {
}
async fn stream_output(
process_id: String,
process_id: ProcessId,
stream: ExecOutputStream,
mut receiver: tokio::sync::mpsc::Receiver<Vec<u8>>,
inner: Arc<Inner>,
output_notify: Arc<Notify>,
) {
while let Some(chunk) = receiver.recv().await {
let _chunk_len = chunk.len();
let notification = {
let mut processes = inner.processes.lock().await;
let Some(entry) = processes.get_mut(&process_id) else {
@@ -448,21 +536,16 @@ async fn stream_output(
break;
};
process.retained_bytes = process.retained_bytes.saturating_sub(evicted.chunk.len());
warn!(
"retained output cap exceeded for process {process_id}; dropping oldest output"
);
}
let _ = process.wake_tx.send(seq);
ExecOutputDeltaNotification {
process_id: process_id.clone(),
seq,
stream,
chunk: chunk.into(),
}
};
output_notify.notify_waiters();
let _ = inner
.events_tx
.send(ExecServerEvent::OutputDelta(notification.clone()));
if inner
.notifications
.notify(crate::protocol::EXEC_OUTPUT_DELTA_METHOD, &notification)
@@ -472,38 +555,46 @@ async fn stream_output(
break;
}
}
finish_output_stream(process_id, inner).await;
}
async fn watch_exit(
process_id: String,
process_id: ProcessId,
exit_rx: tokio::sync::oneshot::Receiver<i32>,
inner: Arc<Inner>,
output_notify: Arc<Notify>,
) {
let exit_code = exit_rx.await.unwrap_or(-1);
{
let notification = {
let mut processes = inner.processes.lock().await;
if let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) {
let seq = process.next_seq;
process.next_seq += 1;
process.exit_code = Some(exit_code);
let _ = process.wake_tx.send(seq);
Some(ExecExitedNotification {
process_id: process_id.clone(),
seq,
exit_code,
})
} else {
None
}
}
output_notify.notify_waiters();
let notification = ExecExitedNotification {
process_id: process_id.clone(),
exit_code,
};
let _ = inner
.events_tx
.send(ExecServerEvent::Exited(notification.clone()));
if inner
.notifications
.notify(crate::protocol::EXEC_EXITED_METHOD, &notification)
.await
.is_err()
output_notify.notify_waiters();
if let Some(notification) = notification
&& inner
.notifications
.notify(crate::protocol::EXEC_EXITED_METHOD, &notification)
.await
.is_err()
{
return;
}
maybe_emit_closed(process_id.clone(), Arc::clone(&inner)).await;
tokio::time::sleep(EXITED_PROCESS_RETENTION).await;
let mut processes = inner.processes.lock().await;
if matches!(
@@ -513,3 +604,51 @@ async fn watch_exit(
processes.remove(&process_id);
}
}
async fn finish_output_stream(process_id: ProcessId, inner: Arc<Inner>) {
{
let mut processes = inner.processes.lock().await;
let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) else {
return;
};
if process.open_streams > 0 {
process.open_streams -= 1;
}
}
maybe_emit_closed(process_id, inner).await;
}
async fn maybe_emit_closed(process_id: ProcessId, inner: Arc<Inner>) {
let notification = {
let mut processes = inner.processes.lock().await;
let Some(ProcessEntry::Running(process)) = processes.get_mut(&process_id) else {
return;
};
if process.closed || process.open_streams != 0 || process.exit_code.is_none() {
return;
}
process.closed = true;
let seq = process.next_seq;
process.next_seq += 1;
let _ = process.wake_tx.send(seq);
Some(ExecClosedNotification {
process_id: process_id.clone(),
seq,
})
};
let Some(notification) = notification else {
return;
};
if inner
.notifications
.notify(EXEC_CLOSED_METHOD, &notification)
.await
.is_err()
{}
}
+20 -18
View File
@@ -1,35 +1,37 @@
use std::sync::Arc;
use async_trait::async_trait;
use tokio::sync::broadcast;
use tokio::sync::watch;
use crate::ExecServerError;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
use crate::ProcessId;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::ReadParams;
use crate::protocol::ReadResponse;
use crate::protocol::TerminateResponse;
use crate::protocol::WriteResponse;
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ExecServerEvent {
OutputDelta(ExecOutputDeltaNotification),
Exited(ExecExitedNotification),
pub struct StartedExecProcess {
pub process: Arc<dyn ExecProcess>,
}
#[async_trait]
pub trait ExecProcess: Send + Sync {
async fn start(&self, params: ExecParams) -> Result<ExecResponse, ExecServerError>;
fn process_id(&self) -> &ProcessId;
async fn read(&self, params: ReadParams) -> Result<ReadResponse, ExecServerError>;
fn subscribe_wake(&self) -> watch::Receiver<u64>;
async fn write(
async fn read(
&self,
process_id: &str,
chunk: Vec<u8>,
) -> Result<WriteResponse, ExecServerError>;
after_seq: Option<u64>,
max_bytes: Option<usize>,
wait_ms: Option<u64>,
) -> Result<ReadResponse, ExecServerError>;
async fn terminate(&self, process_id: &str) -> Result<TerminateResponse, ExecServerError>;
async fn write(&self, chunk: Vec<u8>) -> Result<WriteResponse, ExecServerError>;
fn subscribe_events(&self) -> broadcast::Receiver<ExecServerEvent>;
async fn terminate(&self) -> Result<(), ExecServerError>;
}
#[async_trait]
pub trait ExecBackend: Send + Sync {
async fn start(&self, params: ExecParams) -> Result<StartedExecProcess, ExecServerError>;
}
+74
View File
@@ -0,0 +1,74 @@
use std::borrow::Borrow;
use std::fmt;
use std::ops::Deref;
use serde::Deserialize;
use serde::Serialize;
#[derive(Debug, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ProcessId(String);
impl ProcessId {
pub fn new(value: impl Into<String>) -> Self {
Self(value.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
pub fn into_inner(self) -> String {
self.0
}
}
impl Deref for ProcessId {
type Target = str;
fn deref(&self) -> &Self::Target {
self.as_str()
}
}
impl Borrow<str> for ProcessId {
fn borrow(&self) -> &str {
self.as_str()
}
}
impl AsRef<str> for ProcessId {
fn as_ref(&self) -> &str {
self.as_str()
}
}
impl fmt::Display for ProcessId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
self.0.fmt(f)
}
}
impl From<String> for ProcessId {
fn from(value: String) -> Self {
Self(value)
}
}
impl From<&str> for ProcessId {
fn from(value: &str) -> Self {
Self(value.to_string())
}
}
impl From<&String> for ProcessId {
fn from(value: &String) -> Self {
Self(value.clone())
}
}
impl From<ProcessId> for String {
fn from(value: ProcessId) -> Self {
value.0
}
}
+31 -8
View File
@@ -5,6 +5,8 @@ use base64::engine::general_purpose::STANDARD as BASE64_STANDARD;
use serde::Deserialize;
use serde::Serialize;
use crate::ProcessId;
pub const INITIALIZE_METHOD: &str = "initialize";
pub const INITIALIZED_METHOD: &str = "initialized";
pub const EXEC_METHOD: &str = "process/start";
@@ -13,6 +15,7 @@ pub const EXEC_WRITE_METHOD: &str = "process/write";
pub const EXEC_TERMINATE_METHOD: &str = "process/terminate";
pub const EXEC_OUTPUT_DELTA_METHOD: &str = "process/output";
pub const EXEC_EXITED_METHOD: &str = "process/exited";
pub const EXEC_CLOSED_METHOD: &str = "process/closed";
pub const FS_READ_FILE_METHOD: &str = "fs/readFile";
pub const FS_WRITE_FILE_METHOD: &str = "fs/writeFile";
pub const FS_CREATE_DIRECTORY_METHOD: &str = "fs/createDirectory";
@@ -52,7 +55,7 @@ pub struct InitializeResponse {}
pub struct ExecParams {
/// Client-chosen logical process handle scoped to this connection/session.
/// This is a protocol key, not an OS pid.
pub process_id: String,
pub process_id: ProcessId,
pub argv: Vec<String>,
pub cwd: PathBuf,
pub env: HashMap<String, String>,
@@ -63,13 +66,13 @@ pub struct ExecParams {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecResponse {
pub process_id: String,
pub process_id: ProcessId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ReadParams {
pub process_id: String,
pub process_id: ProcessId,
pub after_seq: Option<u64>,
pub max_bytes: Option<usize>,
pub wait_ms: Option<u64>,
@@ -90,25 +93,36 @@ pub struct ReadResponse {
pub next_seq: u64,
pub exited: bool,
pub exit_code: Option<i32>,
pub closed: bool,
pub failure: Option<String>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WriteParams {
pub process_id: String,
pub process_id: ProcessId,
pub chunk: ByteChunk,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum WriteStatus {
Accepted,
UnknownProcess,
StdinClosed,
Starting,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct WriteResponse {
pub accepted: bool,
pub status: WriteStatus,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct TerminateParams {
pub process_id: String,
pub process_id: ProcessId,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
@@ -128,7 +142,8 @@ pub enum ExecOutputStream {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecOutputDeltaNotification {
pub process_id: String,
pub process_id: ProcessId,
pub seq: u64,
pub stream: ExecOutputStream,
pub chunk: ByteChunk,
}
@@ -136,10 +151,18 @@ pub struct ExecOutputDeltaNotification {
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecExitedNotification {
pub process_id: String,
pub process_id: ProcessId,
pub seq: u64,
pub exit_code: i32,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct ExecClosedNotification {
pub process_id: ProcessId,
pub seq: u64,
}
mod base64_bytes {
use super::BASE64_STANDARD;
use base64::Engine as _;
@@ -10,6 +10,7 @@ use codex_app_server_protocol::FsRemoveParams;
use codex_app_server_protocol::FsWriteFileParams;
use codex_utils_absolute_path::AbsolutePathBuf;
use tokio::io;
use tracing::trace;
use crate::CopyOptions;
use crate::CreateDirectoryOptions;
@@ -22,6 +23,7 @@ use crate::ReadDirectoryEntry;
use crate::RemoveOptions;
const INVALID_REQUEST_ERROR_CODE: i64 = -32600;
const NOT_FOUND_ERROR_CODE: i64 = -32004;
#[derive(Clone)]
pub(crate) struct RemoteFileSystem {
@@ -30,6 +32,7 @@ pub(crate) struct RemoteFileSystem {
impl RemoteFileSystem {
pub(crate) fn new(client: ExecServerClient) -> Self {
trace!("remote fs new");
Self { client }
}
}
@@ -37,6 +40,7 @@ impl RemoteFileSystem {
#[async_trait]
impl ExecutorFileSystem for RemoteFileSystem {
async fn read_file(&self, path: &AbsolutePathBuf) -> FileSystemResult<Vec<u8>> {
trace!("remote fs read_file");
let response = self
.client
.fs_read_file(FsReadFileParams { path: path.clone() })
@@ -51,6 +55,7 @@ impl ExecutorFileSystem for RemoteFileSystem {
}
async fn write_file(&self, path: &AbsolutePathBuf, contents: Vec<u8>) -> FileSystemResult<()> {
trace!("remote fs write_file");
self.client
.fs_write_file(FsWriteFileParams {
path: path.clone(),
@@ -66,6 +71,7 @@ impl ExecutorFileSystem for RemoteFileSystem {
path: &AbsolutePathBuf,
options: CreateDirectoryOptions,
) -> FileSystemResult<()> {
trace!("remote fs create_directory");
self.client
.fs_create_directory(FsCreateDirectoryParams {
path: path.clone(),
@@ -77,6 +83,7 @@ impl ExecutorFileSystem for RemoteFileSystem {
}
async fn get_metadata(&self, path: &AbsolutePathBuf) -> FileSystemResult<FileMetadata> {
trace!("remote fs get_metadata");
let response = self
.client
.fs_get_metadata(FsGetMetadataParams { path: path.clone() })
@@ -94,6 +101,7 @@ impl ExecutorFileSystem for RemoteFileSystem {
&self,
path: &AbsolutePathBuf,
) -> FileSystemResult<Vec<ReadDirectoryEntry>> {
trace!("remote fs read_directory");
let response = self
.client
.fs_read_directory(FsReadDirectoryParams { path: path.clone() })
@@ -111,6 +119,7 @@ impl ExecutorFileSystem for RemoteFileSystem {
}
async fn remove(&self, path: &AbsolutePathBuf, options: RemoveOptions) -> FileSystemResult<()> {
trace!("remote fs remove");
self.client
.fs_remove(FsRemoveParams {
path: path.clone(),
@@ -128,6 +137,7 @@ impl ExecutorFileSystem for RemoteFileSystem {
destination_path: &AbsolutePathBuf,
options: CopyOptions,
) -> FileSystemResult<()> {
trace!("remote fs copy");
self.client
.fs_copy(FsCopyParams {
source_path: source_path.clone(),
@@ -142,6 +152,9 @@ impl ExecutorFileSystem for RemoteFileSystem {
fn map_remote_error(error: ExecServerError) -> io::Error {
match error {
ExecServerError::Server { code, message } if code == NOT_FOUND_ERROR_CODE => {
io::Error::new(io::ErrorKind::NotFound, message)
}
ExecServerError::Server { code, message } if code == INVALID_REQUEST_ERROR_CODE => {
io::Error::new(io::ErrorKind::InvalidInput, message)
}
+64 -29
View File
@@ -1,15 +1,17 @@
use async_trait::async_trait;
use tokio::sync::broadcast;
use std::sync::Arc;
use async_trait::async_trait;
use tokio::sync::watch;
use tracing::trace;
use crate::ExecBackend;
use crate::ExecProcess;
use crate::ExecServerClient;
use crate::ExecServerError;
use crate::ExecServerEvent;
use crate::StartedExecProcess;
use crate::client::ExecServerClient;
use crate::client::Session;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::ReadParams;
use crate::protocol::ReadResponse;
use crate::protocol::TerminateResponse;
use crate::protocol::WriteResponse;
#[derive(Clone)]
@@ -17,35 +19,68 @@ pub(crate) struct RemoteProcess {
client: ExecServerClient,
}
struct RemoteExecProcess {
session: Session,
}
impl RemoteProcess {
pub(crate) fn new(client: ExecServerClient) -> Self {
trace!("remote process new");
Self { client }
}
}
#[async_trait]
impl ExecProcess for RemoteProcess {
async fn start(&self, params: ExecParams) -> Result<ExecResponse, ExecServerError> {
self.client.exec(params).await
}
impl ExecBackend for RemoteProcess {
async fn start(&self, params: ExecParams) -> Result<StartedExecProcess, ExecServerError> {
let process_id = params.process_id.clone();
let session = self.client.register_session(&process_id).await?;
if let Err(err) = self.client.exec(params).await {
session.unregister().await;
return Err(err);
}
async fn read(&self, params: ReadParams) -> Result<ReadResponse, ExecServerError> {
self.client.read(params).await
}
async fn write(
&self,
process_id: &str,
chunk: Vec<u8>,
) -> Result<WriteResponse, ExecServerError> {
self.client.write(process_id, chunk).await
}
async fn terminate(&self, process_id: &str) -> Result<TerminateResponse, ExecServerError> {
self.client.terminate(process_id).await
}
fn subscribe_events(&self) -> broadcast::Receiver<ExecServerEvent> {
self.client.event_receiver()
Ok(StartedExecProcess {
process: Arc::new(RemoteExecProcess { session }),
})
}
}
#[async_trait]
impl ExecProcess for RemoteExecProcess {
fn process_id(&self) -> &crate::ProcessId {
self.session.process_id()
}
fn subscribe_wake(&self) -> watch::Receiver<u64> {
self.session.subscribe_wake()
}
async fn read(
&self,
after_seq: Option<u64>,
max_bytes: Option<usize>,
wait_ms: Option<u64>,
) -> Result<ReadResponse, ExecServerError> {
self.session.read(after_seq, max_bytes, wait_ms).await
}
async fn write(&self, chunk: Vec<u8>) -> Result<WriteResponse, ExecServerError> {
trace!("exec process write");
self.session.write(chunk).await
}
async fn terminate(&self) -> Result<(), ExecServerError> {
trace!("exec process terminate");
self.session.terminate().await
}
}
impl Drop for RemoteExecProcess {
fn drop(&mut self) {
let session = self.session.clone();
tokio::spawn(async move {
session.unregister().await;
});
}
}
+10 -3
View File
@@ -19,7 +19,6 @@ use tokio::sync::Mutex;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::task::JoinHandle;
use tracing::warn;
use crate::connection::JsonRpcConnection;
use crate::connection::JsonRpcConnectionEvent;
@@ -192,12 +191,12 @@ impl RpcClient {
if let Err(err) =
handle_server_message(&pending_for_reader, &event_tx, message).await
{
warn!("JSON-RPC client closing after protocol error: {err}");
let _ = err;
break;
}
}
JsonRpcConnectionEvent::MalformedMessage { reason } => {
warn!("JSON-RPC client closing after malformed message: {reason}");
let _ = reason;
break;
}
JsonRpcConnectionEvent::Disconnected { reason } => {
@@ -357,6 +356,14 @@ pub(crate) fn invalid_params(message: String) -> JSONRPCErrorError {
}
}
pub(crate) fn not_found(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32004,
data: None,
message,
}
}
pub(crate) fn internal_error(message: String) -> JSONRPCErrorError {
JSONRPCErrorError {
code: -32603,
@@ -26,6 +26,7 @@ use crate::RemoveOptions;
use crate::local_file_system::LocalFileSystem;
use crate::rpc::internal_error;
use crate::rpc::invalid_request;
use crate::rpc::not_found;
#[derive(Clone, Default)]
pub(crate) struct FileSystemHandler {
@@ -153,7 +154,9 @@ impl FileSystemHandler {
}
fn map_fs_error(err: io::Error) -> JSONRPCErrorError {
if err.kind() == io::ErrorKind::InvalidInput {
if err.kind() == io::ErrorKind::NotFound {
not_found(err.to_string())
} else if err.kind() == io::ErrorKind::InvalidInput {
invalid_request(err.to_string())
} else {
internal_error(err.to_string())
@@ -6,6 +6,7 @@ use pretty_assertions::assert_eq;
use tokio::sync::mpsc;
use super::ExecServerHandler;
use crate::ProcessId;
use crate::protocol::ExecParams;
use crate::protocol::InitializeResponse;
use crate::protocol::TerminateParams;
@@ -18,7 +19,7 @@ fn exec_params(process_id: &str) -> ExecParams {
env.insert("PATH".to_string(), path.to_string_lossy().into_owned());
}
ExecParams {
process_id: process_id.to_string(),
process_id: ProcessId::from(process_id),
argv: vec![
"bash".to_string(),
"-lc".to_string(),
@@ -84,7 +85,7 @@ async fn terminate_reports_false_after_process_exit() {
loop {
let response = handler
.terminate(TerminateParams {
process_id: "proc-1".to_string(),
process_id: ProcessId::from("proc-1"),
})
.await
.expect("terminate response");
@@ -1,3 +1,5 @@
use std::net::SocketAddr;
use pretty_assertions::assert_eq;
use super::DEFAULT_LISTEN_URL;
@@ -9,7 +11,9 @@ fn parse_listen_url_accepts_default_websocket_url() {
parse_listen_url(DEFAULT_LISTEN_URL).expect("default listen URL should parse");
assert_eq!(
bind_address,
"127.0.0.1:0".parse().expect("valid socket address")
"127.0.0.1:0"
.parse::<SocketAddr>()
.expect("valid socket address")
);
}
@@ -19,7 +23,9 @@ fn parse_listen_url_accepts_websocket_url() {
parse_listen_url("ws://127.0.0.1:1234").expect("websocket listen URL should parse");
assert_eq!(
bind_address,
"127.0.0.1:1234".parse().expect("valid socket address")
"127.0.0.1:1234"
.parse::<SocketAddr>()
.expect("valid socket address")
);
}
+225 -31
View File
@@ -6,19 +6,24 @@ use std::sync::Arc;
use anyhow::Result;
use codex_exec_server::Environment;
use codex_exec_server::ExecBackend;
use codex_exec_server::ExecParams;
use codex_exec_server::ExecProcess;
use codex_exec_server::ExecResponse;
use codex_exec_server::ReadParams;
use codex_exec_server::ProcessId;
use codex_exec_server::ReadResponse;
use codex_exec_server::StartedExecProcess;
use pretty_assertions::assert_eq;
use test_case::test_case;
use tokio::sync::watch;
use tokio::time::Duration;
use tokio::time::timeout;
use common::exec_server::ExecServerHarness;
use common::exec_server::exec_server;
struct ProcessContext {
process: Arc<dyn ExecProcess>,
_server: Option<ExecServerHarness>,
backend: Arc<dyn ExecBackend>,
server: Option<ExecServerHarness>,
}
async fn create_process_context(use_remote: bool) -> Result<ProcessContext> {
@@ -26,24 +31,24 @@ async fn create_process_context(use_remote: bool) -> Result<ProcessContext> {
let server = exec_server().await?;
let environment = Environment::create(Some(server.websocket_url().to_string())).await?;
Ok(ProcessContext {
process: environment.get_executor(),
_server: Some(server),
backend: environment.get_exec_backend(),
server: Some(server),
})
} else {
let environment = Environment::create(None).await?;
let environment = Environment::create(/*exec_server_url*/ None).await?;
Ok(ProcessContext {
process: environment.get_executor(),
_server: None,
backend: environment.get_exec_backend(),
server: None,
})
}
}
async fn assert_exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
let context = create_process_context(use_remote).await?;
let response = context
.process
let session = context
.backend
.start(ExecParams {
process_id: "proc-1".to_string(),
process_id: ProcessId::from("proc-1"),
argv: vec!["true".to_string()],
cwd: std::env::current_dir()?,
env: Default::default(),
@@ -51,30 +56,198 @@ async fn assert_exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
arg0: None,
})
.await?;
assert_eq!(
response,
ExecResponse {
process_id: "proc-1".to_string(),
}
);
assert_eq!(session.process.process_id().as_str(), "proc-1");
let wake_rx = session.process.subscribe_wake();
let (_, exit_code, closed) =
collect_process_output_from_reads(session.process, wake_rx).await?;
let mut next_seq = 0;
assert_eq!(exit_code, Some(0));
assert!(closed);
Ok(())
}
async fn read_process_until_change(
session: Arc<dyn ExecProcess>,
wake_rx: &mut watch::Receiver<u64>,
after_seq: Option<u64>,
) -> Result<ReadResponse> {
let response = session
.read(after_seq, /*max_bytes*/ None, /*wait_ms*/ Some(0))
.await?;
if !response.chunks.is_empty() || response.closed || response.failure.is_some() {
return Ok(response);
}
timeout(Duration::from_secs(2), wake_rx.changed()).await??;
session
.read(after_seq, /*max_bytes*/ None, /*wait_ms*/ Some(0))
.await
.map_err(Into::into)
}
async fn collect_process_output_from_reads(
session: Arc<dyn ExecProcess>,
mut wake_rx: watch::Receiver<u64>,
) -> Result<(String, Option<i32>, bool)> {
let mut output = String::new();
let mut exit_code = None;
let mut after_seq = None;
loop {
let read = context
.process
.read(ReadParams {
process_id: "proc-1".to_string(),
after_seq: Some(next_seq),
max_bytes: None,
wait_ms: Some(100),
})
.await?;
next_seq = read.next_seq;
if read.exited {
assert_eq!(read.exit_code, Some(0));
let response =
read_process_until_change(Arc::clone(&session), &mut wake_rx, after_seq).await?;
if let Some(message) = response.failure {
anyhow::bail!("process failed before closed state: {message}");
}
for chunk in response.chunks {
output.push_str(&String::from_utf8_lossy(&chunk.chunk.into_inner()));
after_seq = Some(chunk.seq);
}
if response.exited {
exit_code = response.exit_code;
}
if response.closed {
break;
}
after_seq = response.next_seq.checked_sub(1).or(after_seq);
}
drop(session);
Ok((output, exit_code, true))
}
async fn assert_exec_process_streams_output(use_remote: bool) -> Result<()> {
let context = create_process_context(use_remote).await?;
let process_id = "proc-stream".to_string();
let session = context
.backend
.start(ExecParams {
process_id: process_id.clone().into(),
argv: vec![
"/bin/sh".to_string(),
"-c".to_string(),
"sleep 0.05; printf 'session output\\n'".to_string(),
],
cwd: std::env::current_dir()?,
env: Default::default(),
tty: false,
arg0: None,
})
.await?;
assert_eq!(session.process.process_id().as_str(), process_id);
let StartedExecProcess { process } = session;
let wake_rx = process.subscribe_wake();
let (output, exit_code, closed) = collect_process_output_from_reads(process, wake_rx).await?;
assert_eq!(output, "session output\n");
assert_eq!(exit_code, Some(0));
assert!(closed);
Ok(())
}
async fn assert_exec_process_write_then_read(use_remote: bool) -> Result<()> {
let context = create_process_context(use_remote).await?;
let process_id = "proc-stdin".to_string();
let session = context
.backend
.start(ExecParams {
process_id: process_id.clone().into(),
argv: vec![
"/usr/bin/python3".to_string(),
"-c".to_string(),
"import sys; line = sys.stdin.readline(); sys.stdout.write(f'from-stdin:{line}'); sys.stdout.flush()".to_string(),
],
cwd: std::env::current_dir()?,
env: Default::default(),
tty: true,
arg0: None,
})
.await?;
assert_eq!(session.process.process_id().as_str(), process_id);
tokio::time::sleep(Duration::from_millis(200)).await;
session.process.write(b"hello\n".to_vec()).await?;
let StartedExecProcess { process } = session;
let wake_rx = process.subscribe_wake();
let (output, exit_code, closed) = collect_process_output_from_reads(process, wake_rx).await?;
assert!(
output.contains("from-stdin:hello"),
"unexpected output: {output:?}"
);
assert_eq!(exit_code, Some(0));
assert!(closed);
Ok(())
}
async fn assert_exec_process_preserves_queued_events_before_subscribe(
use_remote: bool,
) -> Result<()> {
let context = create_process_context(use_remote).await?;
let session = context
.backend
.start(ExecParams {
process_id: ProcessId::from("proc-queued"),
argv: vec![
"/bin/sh".to_string(),
"-c".to_string(),
"printf 'queued output\\n'".to_string(),
],
cwd: std::env::current_dir()?,
env: Default::default(),
tty: false,
arg0: None,
})
.await?;
tokio::time::sleep(Duration::from_millis(200)).await;
let StartedExecProcess { process } = session;
let wake_rx = process.subscribe_wake();
let (output, exit_code, closed) = collect_process_output_from_reads(process, wake_rx).await?;
assert_eq!(output, "queued output\n");
assert_eq!(exit_code, Some(0));
assert!(closed);
Ok(())
}
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn remote_exec_process_reports_transport_disconnect() -> Result<()> {
let mut context = create_process_context(/*use_remote*/ true).await?;
let session = context
.backend
.start(ExecParams {
process_id: ProcessId::from("proc-disconnect"),
argv: vec![
"/bin/sh".to_string(),
"-c".to_string(),
"sleep 10".to_string(),
],
cwd: std::env::current_dir()?,
env: Default::default(),
tty: false,
arg0: None,
})
.await?;
let server = context
.server
.as_mut()
.expect("remote context should include exec-server harness");
server.shutdown().await?;
let mut wake_rx = session.process.subscribe_wake();
let response =
read_process_until_change(session.process, &mut wake_rx, /*after_seq*/ None).await?;
let message = response
.failure
.expect("disconnect should surface as a failure");
assert!(
message.starts_with("exec-server transport disconnected"),
"unexpected failure message: {message}"
);
assert!(
response.closed,
"disconnect should close the process session"
);
Ok(())
}
@@ -85,3 +258,24 @@ async fn assert_exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
async fn exec_process_starts_and_exits(use_remote: bool) -> Result<()> {
assert_exec_process_starts_and_exits(use_remote).await
}
#[test_case(false ; "local")]
#[test_case(true ; "remote")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_process_streams_output(use_remote: bool) -> Result<()> {
assert_exec_process_streams_output(use_remote).await
}
#[test_case(false ; "local")]
#[test_case(true ; "remote")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_process_write_then_read(use_remote: bool) -> Result<()> {
assert_exec_process_write_then_read(use_remote).await
}
#[test_case(false ; "local")]
#[test_case(true ; "remote")]
#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_process_preserves_queued_events_before_subscribe(use_remote: bool) -> Result<()> {
assert_exec_process_preserves_queued_events_before_subscribe(use_remote).await
}
+7 -1
View File
@@ -36,7 +36,7 @@ async fn create_file_system_context(use_remote: bool) -> Result<FileSystemContex
_server: Some(server),
})
} else {
let environment = Environment::create(None).await?;
let environment = Environment::create(/*exec_server_url*/ None).await?;
Ok(FileSystemContext {
file_system: environment.get_filesystem(),
_server: None,
@@ -122,6 +122,12 @@ async fn file_system_methods_cover_surface_area(use_remote: bool) -> Result<()>
.with_context(|| format!("mode={use_remote}"))?;
assert_eq!(nested_file_contents, b"hello from trait");
let nested_file_text = file_system
.read_file_text(&absolute_path(nested_file.clone()))
.await
.with_context(|| format!("mode={use_remote}"))?;
assert_eq!(nested_file_text, "hello from trait");
file_system
.copy(
&absolute_path(nested_file),
+2 -1
View File
@@ -6,6 +6,7 @@ use codex_app_server_protocol::JSONRPCMessage;
use codex_app_server_protocol::JSONRPCResponse;
use codex_exec_server::ExecResponse;
use codex_exec_server::InitializeParams;
use codex_exec_server::ProcessId;
use common::exec_server::exec_server;
use pretty_assertions::assert_eq;
@@ -62,7 +63,7 @@ async fn exec_server_starts_process_over_websocket() -> anyhow::Result<()> {
assert_eq!(
process_start_response,
ExecResponse {
process_id: "proc-1".to_string()
process_id: ProcessId::from("proc-1")
}
);