Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
20 changes: 20 additions & 0 deletions codex-rs/exec-server-protocol/src/protocol.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@ 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 ENVIRONMENT_INFO_METHOD: &str = "environment/info";
pub const ENVIRONMENT_STATUS_METHOD: &str = "environment/status";
pub const FS_READ_FILE_METHOD: &str = "fs/readFile";
pub const FS_OPEN_METHOD: &str = "fs/open";
pub const FS_READ_BLOCK_METHOD: &str = "fs/readBlock";
Expand Down Expand Up @@ -83,6 +84,25 @@ pub struct EnvironmentInfo {
pub cwd: Option<PathUri>,
}

/// Status returned by an initialized exec-server connection.
///
/// The response is intentionally small today. New status details can be added
/// without changing the method used by clients to verify that an initialized
/// exec-server connection is still responsive.
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub struct EnvironmentStatus {
pub status: EnvironmentStatusKind,
}

/// High-level status reported by exec-server itself.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "camelCase")]
pub enum EnvironmentStatusKind {
/// The connection is initialized and exec-server can handle requests.
Ready,
}

impl EnvironmentInfo {
/// Returns information about the current local exec-server process.
pub fn local() -> Self {
Expand Down
42 changes: 42 additions & 0 deletions codex-rs/exec-server/src/client.rs
Original file line number Diff line number Diff line change
Expand Up @@ -39,6 +39,7 @@ use crate::process::ExecProcessEvent;
use crate::process::ExecProcessEventLog;
use crate::process::ExecProcessEventReceiver;
use crate::protocol::ENVIRONMENT_INFO_METHOD;
use crate::protocol::ENVIRONMENT_STATUS_METHOD;
use crate::protocol::EXEC_CLOSED_METHOD;
use crate::protocol::EXEC_EXITED_METHOD;
use crate::protocol::EXEC_METHOD;
Expand All @@ -48,6 +49,7 @@ use crate::protocol::EXEC_SIGNAL_METHOD;
use crate::protocol::EXEC_TERMINATE_METHOD;
use crate::protocol::EXEC_WRITE_METHOD;
use crate::protocol::EnvironmentInfo;
use crate::protocol::EnvironmentStatus;
use crate::protocol::ExecClosedNotification;
use crate::protocol::ExecExitedNotification;
use crate::protocol::ExecOutputDeltaNotification;
Expand Down Expand Up @@ -114,6 +116,7 @@ mod recovery;
const CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
const INITIALIZE_TIMEOUT: Duration = Duration::from_secs(10);
const ENVIRONMENT_INFO_TIMEOUT: Duration = Duration::from_secs(30);
const ENVIRONMENT_STATUS_TIMEOUT: Duration = Duration::from_secs(10);
const PROCESS_EVENT_CHANNEL_CAPACITY: usize = 256;
const PROCESS_EVENT_RETAINED_BYTES: usize = 1024 * 1024;
const MAX_PENDING_PROCESS_EVENTS: usize = 256;
Expand Down Expand Up @@ -302,6 +305,30 @@ impl LazyRemoteExecServerClient {
})
}

pub(crate) async fn status(&self) -> crate::EnvironmentObservedStatus {
// Fail-fast lookup preserves the non-mutating contract: never start or recover a client.
let client = match self.fail_fast().get().await {
Ok(client) => client,
Err(error) => {
// Without a completed startup attempt, there is no exec-server connection to probe.
if self.cached_client().is_none() && self.startup.get().is_none() {
return crate::EnvironmentObservedStatus::Pending;
}
// A known connection failure is reported without retrying it as part of status.
return crate::EnvironmentObservedStatus::Disconnected {
error: error.to_string(),
};
}
};
// Every callable client is probed so callers never receive a cached health result.
match client.environment_status().await {
Ok(_) => crate::EnvironmentObservedStatus::Ready,
Err(error) => crate::EnvironmentObservedStatus::Disconnected {
error: error.to_string(),
},
}
}

pub(crate) fn fail_fast(&self) -> Self {
Self {
recovery_policy: RecoveryPolicy::FailFast,
Expand Down Expand Up @@ -604,6 +631,16 @@ impl ExecServerClient {
)
}

pub async fn environment_status(&self) -> Result<EnvironmentStatus, ExecServerError> {
// Health checks only reuse an existing RPC connection and never initiate recovery.
let rpc_client = self.rpc_client_without_recovery()?;
map_rpc_call_result(
rpc_client
.call_with_timeout(ENVIRONMENT_STATUS_METHOD, &(), ENVIRONMENT_STATUS_TIMEOUT)
.await,
)
}

pub async fn read(&self, params: ReadParams) -> Result<ReadResponse, ExecServerError> {
self.call(EXEC_READ_METHOD, &params).await
}
Expand Down Expand Up @@ -1432,6 +1469,7 @@ mod tests {
use super::ExecServerClient;
use super::ExecServerClientConnectOptions;
use super::LazyRemoteExecServerClient;
use crate::EnvironmentObservedStatus;
use crate::ProcessId;
#[cfg(not(windows))]
use crate::client_api::DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT;
Expand Down Expand Up @@ -2411,6 +2449,10 @@ mod tests {
Ok(_) => panic!("burned environment should stay failed"),
Err(error) => error,
};
assert!(matches!(
client.status().await,
EnvironmentObservedStatus::Disconnected { .. }
));

let (
super::ExecServerError::ConnectionAttempt(first),
Expand Down
82 changes: 82 additions & 0 deletions codex-rs/exec-server/src/environment.rs
Original file line number Diff line number Diff line change
Expand Up @@ -69,6 +69,30 @@ pub struct PendingEnvironmentRegistration(oneshot::Sender<Result<String, String>
pub const LOCAL_ENVIRONMENT_ID: &str = "local";
pub const REMOTE_ENVIRONMENT_ID: &str = "remote";

/// Non-mutating connection status observed by an environment owner.
///
/// Computing this status never starts, waits for, or reconnects an exec-server
/// transport. Already-ready remote environments may receive a fail-fast probe
/// over their existing connection.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum EnvironmentObservedStatus {
/// A local environment, or a remote environment whose existing connection answered a probe.
Ready,
/// The configured environment has no ready connection and no observed connection failure.
///
/// This includes lazy transports that have never been started and initial startup that has
/// not finished. Computing status does not start the environment or wait for startup.
Pending,
/// A connection attempt, prior connection, or fail-fast status probe observed a failure.
///
/// This does not promise that the failure is terminal: later normal environment use may
/// recover the connection. Computing status itself does not trigger recovery.
Disconnected {
/// Human-readable reason recorded by the failed connection attempt or probe.
error: String,
},
}

impl EnvironmentManager {
/// Builds a test-only manager without configured sandbox helper paths.
pub fn default_for_tests() -> Self {
Expand Down Expand Up @@ -273,6 +297,19 @@ impl EnvironmentManager {
environment_ids
}

/// Returns every configured environment id in deterministic order.
pub fn environment_ids(&self) -> Vec<String> {
let mut environment_ids = self
.environments
.read()
.unwrap_or_else(std::sync::PoisonError::into_inner)
.keys()
.cloned()
.collect::<Vec<_>>();
environment_ids.sort_unstable();
environment_ids
}

/// Returns the local environment instance when one is configured.
pub fn try_local_environment(&self) -> Option<Arc<Environment>> {
self.local_environment.as_ref().map(Arc::clone)
Expand All @@ -293,6 +330,15 @@ impl EnvironmentManager {
.cloned()
}

/// Returns the current status of one named environment when it is configured.
pub async fn get_environment_status(
&self,
environment_id: &str,
) -> Option<EnvironmentObservedStatus> {
let environment = self.get_environment(environment_id)?;
Some(environment.status().await)
}

/// Adds or replaces a named remote environment without changing the
/// manager's default environment selection. Uses the default WebSocket
/// connection timeout when none is provided.
Expand Down Expand Up @@ -645,6 +691,19 @@ impl Environment {
}
}

/// Returns the environment's status without starting or recovering it.
///
/// Local environments are always ready. Remote environments with an
/// already-ready cached connection receive a fail-fast `environment/status`
/// probe; other remote states are returned from cached connection state
/// without waiting for startup or recovery.
pub async fn status(&self) -> EnvironmentObservedStatus {
match &self.remote_client {
Some(client) => client.status().await,
None => EnvironmentObservedStatus::Ready,
}
}

pub fn get_exec_backend(&self) -> Arc<dyn ExecBackend> {
Arc::clone(&self.exec_backend)
}
Expand Down Expand Up @@ -674,6 +733,7 @@ mod tests {

use super::Environment;
use super::EnvironmentManager;
use super::EnvironmentObservedStatus;
use super::LOCAL_ENVIRONMENT_ID;
use super::REMOTE_ENVIRONMENT_ID;
use super::noise_environment_config_from_values;
Expand Down Expand Up @@ -1094,6 +1154,28 @@ mod tests {
.expect("accept connection");
}

#[tokio::test]
async fn environment_status_keeps_stdio_environment_pending() {
let environment = Environment::remote_with_transport(
ExecServerTransportParams::StdioCommand {
command: StdioExecServerCommand {
program: "codex-missing-exec-server-for-test".to_string(),
args: Vec::new(),
env: HashMap::new(),
cwd: None,
},
initialize_timeout: Duration::from_secs(1),
},
/*local_runtime_paths*/ None,
);

assert_eq!(
environment.status().await,
EnvironmentObservedStatus::Pending
);
assert!(!environment.startup_finished());
}

#[tokio::test]
async fn environment_manager_leaves_stdio_environment_lazy() {
let environment = Environment::remote_with_transport(
Expand Down
3 changes: 3 additions & 0 deletions codex-rs/exec-server/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -67,6 +67,7 @@ pub use environment::CODEX_EXEC_SERVER_NOISE_REGISTRY_URL_ENV_VAR;
pub use environment::CODEX_EXEC_SERVER_URL_ENV_VAR;
pub use environment::Environment;
pub use environment::EnvironmentManager;
pub use environment::EnvironmentObservedStatus;
pub use environment::LOCAL_ENVIRONMENT_ID;
pub use environment::PendingEnvironmentRegistration;
pub use environment::REMOTE_ENVIRONMENT_ID;
Expand Down Expand Up @@ -95,6 +96,8 @@ pub use process::ExecProcessFuture;
pub use process::StartedExecProcess;
pub use protocol::ByteChunk;
pub use protocol::EnvironmentInfo;
pub use protocol::EnvironmentStatus;
pub use protocol::EnvironmentStatusKind;
pub use protocol::ExecClosedNotification;
pub use protocol::ExecEnvPolicy;
pub use protocol::ExecExitedNotification;
Expand Down
9 changes: 9 additions & 0 deletions codex-rs/exec-server/src/server/handler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,8 @@ use crate::ExecServerRuntimePaths;
use crate::client::http_client::PendingReqwestHttpBodyStream;
use crate::client::http_client::ReqwestHttpRequestRunner;
use crate::protocol::EnvironmentInfo;
use crate::protocol::EnvironmentStatus;
use crate::protocol::EnvironmentStatusKind;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::FsCanonicalizeParams;
Expand Down Expand Up @@ -166,6 +168,13 @@ impl ExecServerHandler {
Ok(EnvironmentInfo::local())
}

pub(crate) fn environment_status(&self) -> Result<EnvironmentStatus, JSONRPCErrorError> {
self.require_initialized_for("environment status")?;
Ok(EnvironmentStatus {
status: EnvironmentStatusKind::Ready,
})
}

pub(crate) async fn exec_read(
&self,
params: ReadParams,
Expand Down
10 changes: 10 additions & 0 deletions codex-rs/exec-server/src/server/processor.rs
Original file line number Diff line number Diff line change
Expand Up @@ -308,10 +308,13 @@ mod tests {
use crate::ProcessId;
use crate::connection::JsonRpcConnection;
use crate::protocol::ENVIRONMENT_INFO_METHOD;
use crate::protocol::ENVIRONMENT_STATUS_METHOD;
use crate::protocol::EXEC_METHOD;
use crate::protocol::EXEC_READ_METHOD;
use crate::protocol::EXEC_TERMINATE_METHOD;
use crate::protocol::EnvironmentInfo;
use crate::protocol::EnvironmentStatus;
use crate::protocol::EnvironmentStatusKind;
use crate::protocol::ExecParams;
use crate::protocol::ExecResponse;
use crate::protocol::INITIALIZE_METHOD;
Expand Down Expand Up @@ -394,9 +397,16 @@ mod tests {

send_request(&mut writer, /*id*/ 2, ENVIRONMENT_INFO_METHOD, &()).await;
send_request(&mut writer, /*id*/ 3, ENVIRONMENT_INFO_METHOD, &()).await;
send_request(&mut writer, /*id*/ 4, ENVIRONMENT_STATUS_METHOD, &()).await;

let _: EnvironmentInfo = read_response(&mut lines, /*expected_id*/ 2).await;
let _: EnvironmentInfo = read_response(&mut lines, /*expected_id*/ 3).await;
assert_eq!(
read_response::<EnvironmentStatus>(&mut lines, /*expected_id*/ 4).await,
EnvironmentStatus {
status: EnvironmentStatusKind::Ready,
}
);

drop(writer);
drop(lines);
Expand Down
5 changes: 5 additions & 0 deletions codex-rs/exec-server/src/server/registry.rs
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
use std::sync::Arc;

use crate::protocol::ENVIRONMENT_INFO_METHOD;
use crate::protocol::ENVIRONMENT_STATUS_METHOD;
use crate::protocol::EXEC_METHOD;
use crate::protocol::EXEC_READ_METHOD;
use crate::protocol::EXEC_SIGNAL_METHOD;
Expand Down Expand Up @@ -71,6 +72,10 @@ pub(crate) fn build_router() -> RpcRouter<ExecServerHandler> {
ENVIRONMENT_INFO_METHOD,
|handler: Arc<ExecServerHandler>, _params: ()| async move { handler.environment_info() },
);
router.request(
ENVIRONMENT_STATUS_METHOD,
|handler: Arc<ExecServerHandler>, _params: ()| async move { handler.environment_status() },
);
router.request(
EXEC_READ_METHOD,
|handler: Arc<ExecServerHandler>, params: ReadParams| async move {
Expand Down
50 changes: 50 additions & 0 deletions codex-rs/exec-server/tests/health.rs
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,12 @@
mod common;

use codex_exec_server::Environment;
use codex_exec_server::EnvironmentStatus;
use codex_exec_server::EnvironmentStatusKind;
use codex_exec_server::InitializeParams;
use codex_exec_server::InitializeResponse;
use codex_exec_server_protocol::JSONRPCMessage;
use codex_exec_server_protocol::JSONRPCResponse;
use common::exec_server::exec_server;
use pretty_assertions::assert_eq;

Expand Down Expand Up @@ -34,3 +40,47 @@ async fn remote_environment_fetches_info_from_exec_server() -> anyhow::Result<()
server.shutdown().await?;
Ok(())
}

#[tokio::test(flavor = "multi_thread", worker_threads = 2)]
async fn exec_server_reports_environment_status_over_websocket() -> anyhow::Result<()> {
let mut server = exec_server().await?;
let initialize_id = server
.send_request(
"initialize",
serde_json::to_value(InitializeParams {
client_name: "exec-server-health-test".to_string(),
resume_session_id: None,
})?,
)
.await?;
let JSONRPCMessage::Response(JSONRPCResponse {
id,
result: initialize_result,
}) = server.next_event().await?
else {
panic!("expected initialize response");
};
assert_eq!(id, initialize_id);
let _: InitializeResponse = serde_json::from_value(initialize_result)?;
server
.send_notification("initialized", serde_json::json!({}))
.await?;

let status_id = server
.send_request("environment/status", serde_json::json!({}))
.await?;
let JSONRPCMessage::Response(JSONRPCResponse { id, result }) = server.next_event().await?
else {
panic!("expected environment status response");
};
assert_eq!(id, status_id);
assert_eq!(
serde_json::from_value::<EnvironmentStatus>(result)?,
EnvironmentStatus {
status: EnvironmentStatusKind::Ready,
}
);

server.shutdown().await?;
Ok(())
}
Loading