diff --git "a/codex-rs/exec-server/src/client.rs" "b/codex-rs/exec-server/src/client.rs" new file mode 100644--- /dev/null +++ "b/codex-rs/exec-server/src/client.rs" @@ -0,0 +1,3232 @@ +use std::collections::BTreeMap; +use std::collections::HashMap; +use std::sync::Arc; +use std::sync::Mutex as StdMutex; +use std::sync::OnceLock; +use std::sync::atomic::AtomicBool; +use std::sync::atomic::AtomicU64; +use std::sync::atomic::Ordering; +use std::time::Duration; + +use arc_swap::ArcSwap; +use arc_swap::ArcSwapOption; +use codex_exec_server_protocol::JSONRPCNotification; +use codex_network_proxy::NetworkPolicyDecider; +use codex_network_proxy::NetworkProxyAuditMetadata; +use codex_network_proxy::NetworkRequestCancellation; +use codex_network_proxy::NetworkRequestCancellationReason; +use futures::FutureExt; +use futures::future::BoxFuture; +use serde_json::Value; +use tokio::sync::Mutex; +use tokio::sync::OnceCell; +use tokio::sync::Semaphore; +use tokio::sync::mpsc; +use tokio::sync::watch; +use tokio_util::sync::CancellationToken; +use tokio_util::task::AbortOnDropHandle; + +use tokio::time::timeout; +use tracing::Instrument; +use tracing::debug; +use tracing::instrument::WithSubscriber; + +use crate::ProcessId; +use crate::client::http_client::response_body_stream::MAX_QUEUED_HTTP_BODY_BYTES; +use crate::client::http_client::response_body_stream::QueuedHttpBodyDelta; +use crate::client_api::ExecServerClientConnectOptions; +use crate::client_api::ExecServerTransportParams; +use crate::client_api::HttpClient; +use crate::client_api::RemoteExecServerConnectArgs; +use crate::client_api::StdioExecServerConnectArgs; +use crate::client_transport::ExecServerReconnectStrategy; +use crate::connection::JsonRpcConnection; +use crate::environment::EnvironmentConnectionState; +use crate::process::ExecProcessEvent; +use crate::process::ExecProcessEventLog; +use crate::process::ExecProcessEventReceiver; +use crate::protocol::CAPABILITY_ROOTS_DISCOVER_METHOD; +use crate::protocol::CapabilityRootsDiscoverParams; +use crate::protocol::CapabilityRootsDiscoverResponse; +use crate::protocol::ENVIRONMENT_CONFIG_READ_METHOD; +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; +use crate::protocol::EXEC_OUTPUT_DELTA_METHOD; +use crate::protocol::EXEC_READ_METHOD; +use crate::protocol::EXEC_SIGNAL_METHOD; +use crate::protocol::EXEC_TERMINATE_METHOD; +use crate::protocol::EXEC_WRITE_METHOD; +use crate::protocol::EnvironmentConfigReadParams; +use crate::protocol::EnvironmentConfigReadResponse; +use crate::protocol::EnvironmentInfo; +use crate::protocol::EnvironmentStatus; +use crate::protocol::ExecClosedNotification; +use crate::protocol::ExecExitedNotification; +use crate::protocol::ExecOutputDeltaNotification; +use crate::protocol::ExecParams; +use crate::protocol::ExecResponse; +use crate::protocol::FS_CANONICALIZE_METHOD; +use crate::protocol::FS_CLOSE_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_OPEN_METHOD; +use crate::protocol::FS_READ_BLOCK_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_WALK_METHOD; +use crate::protocol::FS_WRITE_FILE_METHOD; +use crate::protocol::FsCanonicalizeParams; +use crate::protocol::FsCanonicalizeResponse; +use crate::protocol::FsCloseParams; +use crate::protocol::FsCloseResponse; +use crate::protocol::FsCopyParams; +use crate::protocol::FsCopyResponse; +use crate::protocol::FsCreateDirectoryParams; +use crate::protocol::FsCreateDirectoryResponse; +use crate::protocol::FsGetMetadataParams; +use crate::protocol::FsGetMetadataResponse; +use crate::protocol::FsOpenParams; +use crate::protocol::FsOpenResponse; +use crate::protocol::FsReadBlockParams; +use crate::protocol::FsReadBlockResponse; +use crate::protocol::FsReadDirectoryParams; +use crate::protocol::FsReadDirectoryResponse; +use crate::protocol::FsReadFileParams; +use crate::protocol::FsReadFileResponse; +use crate::protocol::FsRemoveParams; +use crate::protocol::FsRemoveResponse; +use crate::protocol::FsWalkParams; +use crate::protocol::FsWalkResponse; +use crate::protocol::FsWriteFileParams; +use crate::protocol::FsWriteFileResponse; +use crate::protocol::HTTP_REQUEST_BODY_DELTA_METHOD; +use crate::protocol::INITIALIZE_METHOD; +use crate::protocol::INITIALIZED_METHOD; +use crate::protocol::InitializeParams; +use crate::protocol::InitializeResponse; +use crate::protocol::NETWORK_POLICY_DECISION_METHOD; +use crate::protocol::NetworkPolicyDecisionNotification; +use crate::protocol::ProcessOutputChunk; +use crate::protocol::ProcessSandboxType; +use crate::protocol::ProcessSignal; +use crate::protocol::ReadParams; +use crate::protocol::ReadResponse; +use crate::protocol::SignalParams; +use crate::protocol::SignalResponse; +use crate::protocol::TerminateParams; +use crate::protocol::TerminateResponse; +use crate::protocol::WriteParams; +use crate::protocol::WriteResponse; +use crate::rpc::RpcCallError; +use crate::rpc::RpcClient; +use crate::rpc_server_requests::MAX_IN_FLIGHT_SERVER_CALLS; +use codex_http_client::HttpClientFactory; + +#[path = "client/accepted.rs"] +pub(crate) mod accepted; +pub(crate) mod http_client; +mod network_policy_audit; +#[cfg(test)] +#[path = "../tests/unit/client_provisioning_tests.rs"] +mod provisioning_tests; +#[path = "client_recovery.rs"] +mod recovery; +#[path = "client_refresh.rs"] +mod refresh; +#[cfg(test)] +pub(crate) use recovery::is_environment_offline_error; +pub(crate) use recovery::is_retryable_recovery_error; +pub(crate) use recovery::is_retryable_registry_error; +pub(crate) use recovery::registry_recovery_retry_delay; +use refresh::ConnectionAttempt; + +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; +const MAX_PENDING_PROCESS_EVENT_BYTES: usize = 1024 * 1024; + +impl Default for ExecServerClientConnectOptions { + fn default() -> Self { + Self { + client_name: "codex-core".to_string(), + initialize_timeout: INITIALIZE_TIMEOUT, + resume_session_id: None, + } + } +} + +impl From for ExecServerClientConnectOptions { + fn from(value: RemoteExecServerConnectArgs) -> Self { + Self { + client_name: value.client_name, + initialize_timeout: value.initialize_timeout, + resume_session_id: value.resume_session_id, + } + } +} + +impl From for ExecServerClientConnectOptions { + fn from(value: StdioExecServerConnectArgs) -> Self { + Self { + client_name: value.client_name, + initialize_timeout: value.initialize_timeout, + resume_session_id: value.resume_session_id, + } + } +} + +impl RemoteExecServerConnectArgs { + pub fn new( + websocket_url: String, + client_name: String, + http_client_factory: HttpClientFactory, + ) -> Self { + Self { + websocket_url, + client_name, + connect_timeout: CONNECT_TIMEOUT, + initialize_timeout: INITIALIZE_TIMEOUT, + resume_session_id: None, + http_client_factory, + } + } +} + +pub(crate) struct SessionState { + wake_tx: watch::Sender, + events: ExecProcessEventLog, + ordered_events: StdMutex, + recoverable: AtomicBool, + next_write_id: AtomicU64, + network_policy: NetworkPolicyState, +} + +struct NetworkPolicyState { + controller: ArcSwapOption, + cancelled: CancellationToken, + cancellation: NetworkRequestCancellation, + audit: Option, +} + +#[derive(Clone)] +struct NetworkPolicyDecisionController { + decider: Arc, + timeout: Duration, +} + +struct NetworkPolicyAuditContext { + metadata: NetworkProxyAuditMetadata, + execution_id: Option, +} + +#[derive(Default)] +struct OrderedSessionEvents { + last_published_seq: u64, + exit_published: bool, + closed_published: bool, + // Server-side output, exit, and closed notifications are emitted by + // different tasks and can reach the client out of order. Keep future events + // here until all lower sequence numbers have been published. + pending: BTreeMap, + pending_bytes: usize, + failure: Option, +} + +#[derive(Clone)] +pub(crate) struct Session { + client: ExecServerClient, + process_id: ProcessId, + sandbox_type: Option, + state: Arc, +} + +struct Inner { + connection: StdMutex, + connection_changed: watch::Sender<()>, + // 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>>, + // 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: StdMutex<()>, + // Streaming HTTP responses are keyed by a client-generated request id + // because they share the same connection-global notification channel as + // process output. Keep the routing table local to the client so higher + // layers can consume body chunks like a normal byte stream. + http_body_streams: ArcSwap>>, + http_body_stream_failures: ArcSwap>, + http_body_streams_write_lock: Mutex<()>, + http_body_stream_byte_budget: Arc, + http_body_stream_next_id: AtomicU64, + // Keep admission shared while recovered transports finish older requests. + rpc_inbound_request_slots: Arc, + session_id: OnceLock, + retired: CancellationToken, + /// Caches metadata from initialization or the first successful info request for this client's lifetime. + environment_info: OnceCell, + reconnect_strategy: Option, +} + +struct ConnectionState { + status: ConnectionStatus, + active_process_starts: usize, + environment_connection_state_tx: watch::Sender, +} + +enum ConnectionStatus { + Connected(Arc), + Recovering, + Failed(String), +} + +impl ConnectionState { + fn set_status(&mut self, status: ConnectionStatus) { + self.status = status; + self.publish_environment_connection_state(); + } + + fn publish_environment_connection_state(&self) { + let state = match &self.status { + ConnectionStatus::Connected(rpc_client) if !rpc_client.is_disconnected() => { + EnvironmentConnectionState::Connected + } + ConnectionStatus::Connected(_) + | ConnectionStatus::Recovering + | ConnectionStatus::Failed(_) => EnvironmentConnectionState::Disconnected, + }; + let _ = self + .environment_connection_state_tx + .send_if_modified(|current| { + if *current == state { + false + } else { + *current = state; + true + } + }); + } +} + +#[derive(Clone, Copy)] +enum RecoveryPolicy { + Wait, + FailFast, +} + +#[derive(Clone)] +pub struct ExecServerClient { + inner: Arc, + recovery_policy: RecoveryPolicy, +} + +/// State carried from a Noise readiness wait into the initialize RPC. +/// +/// The span preserves the existing logical initialize operation while +/// `timeout_for_error` keeps diagnostics tied to the caller's configured +/// budget after readiness has consumed part of that budget. +pub(crate) struct NoiseInitializeContext { + pub(crate) span: tracing::Span, + pub(crate) timeout_for_error: Duration, +} + +struct ActiveProcessStart { + inner: Arc, +} + +impl Drop for ActiveProcessStart { + fn drop(&mut self) { + self.inner.finish_process_start(); + } +} + +struct PendingProcessStartSession { + inner: Arc, + process_id: ProcessId, + state: Arc, + armed: bool, +} + +impl Drop for PendingProcessStartSession { + fn drop(&mut self) { + if self.armed { + self.inner.remove_session_if(&self.process_id, &self.state); + } + } +} + +type ConnectionResult = Result>; + +#[derive(Clone)] +pub(crate) struct LazyRemoteExecServerClient { + transport_params: Option, + http_client_factory: HttpClientFactory, + recovery_policy: RecoveryPolicy, + // Saves the first startup result so callers share it; retryable failures use reconnect. + startup: Arc, + // The latest successful client, replaced whenever reconnecting succeeds. + current_client: Arc>>, + reconnect: Arc>>>, + refresh_lock: Arc>, + environment_connection_state_tx: watch::Sender, +} + +impl LazyRemoteExecServerClient { + pub(crate) fn new( + transport_params: ExecServerTransportParams, + http_client_factory: HttpClientFactory, + ) -> Self { + Self { + transport_params: Some(transport_params), + http_client_factory, + recovery_policy: RecoveryPolicy::Wait, + startup: Arc::new(ConnectionAttempt::default()), + current_client: Arc::new(StdMutex::new(None)), + reconnect: Arc::new(StdMutex::new(None)), + refresh_lock: Arc::new(Mutex::new(())), + environment_connection_state_tx: watch::channel( + EnvironmentConnectionState::Disconnected, + ) + .0, + } + } + + pub(crate) fn subscribe_connection_state(&self) -> watch::Receiver { + self.environment_connection_state_tx.subscribe() + } + + pub(crate) fn start_connecting(&self) -> Option> { + // Stdio starts a process, so keep it lazy until the environment is used. + if matches!( + self.transport_params, + Some(ExecServerTransportParams::StdioCommand { .. }) + ) { + return None; + } + let client = self.clone(); + Some(AbortOnDropHandle::new(tokio::spawn( + async move { + if let Err(error) = client.wait_until_ready().await { + debug!(%error, "exec-server environment startup failed"); + } + } + .in_current_span() + .with_current_subscriber(), + ))) + } + + pub(crate) fn startup_finished(&self) -> bool { + // Explicit refresh can install the first client without polling startup. + self.cached_client().is_some() || self.startup.result.get().is_some() + } + + pub(crate) fn readiness_result(&self) -> Option> { + if let Some(client) = self.cached_client() { + return client.readiness_result(); + } + self.startup.result.get().and_then(|result| match result { + Ok(client) => client.readiness_result(), + Err(error) => Some(Err(ExecServerError::ConnectionAttempt(Arc::clone(error)))), + }) + } + + pub(crate) async fn status(&self) -> crate::EnvironmentObservedStatus { + if let Some(ExecServerTransportParams::Deferred(deferred)) = &self.transport_params + && let Some(Err(error)) = deferred.readiness.borrow().as_ref() + { + return crate::EnvironmentObservedStatus::Disconnected { + error: ExecServerError::ProvisioningFailed(error.clone()).to_string(), + }; + } + // 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.result.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, + ..self.clone() + } + } + + pub(crate) async fn wait_until_ready(&self) -> Result<(), ExecServerError> { + self.initial_client().await.map(drop) + } + + pub(crate) async fn get(&self) -> Result { + if matches!(self.recovery_policy, RecoveryPolicy::FailFast) { + let client = match self.cached_client() { + Some(client) => client, + None => match self.startup.result.get() { + Some(Ok(client)) => client.clone(), + Some(Err(error)) => { + return Err(ExecServerError::ConnectionAttempt(Arc::clone(error))); + } + None => { + return Err(ExecServerError::Disconnected( + "exec-server environment is not ready".to_string(), + )); + } + }, + }; + return client.fail_fast(); + } + if let Some(client) = self.connected_client() { + return Ok(client); + } + + let Some(cached_client) = self.cached_client() else { + let client = self.initial_client().await?; + if !client.is_disconnected() || !self.can_reconnect() { + return Ok(client); + } + return self.reconnect().await; + }; + + if !self.can_reconnect() { + return Ok(cached_client); + } + + self.reconnect().await + } + + async fn initial_client(&self) -> Result { + let result = if self.can_reconnect() + && (self.startup.cancelled.is_cancelled() + || self.startup.result.get().is_some_and(|result| { + result + .as_ref() + .is_err_and(|error| recovery::is_retryable_recovery_error(error)) + })) { + Box::pin(self.reconnect()).await + } else { + self.startup + .result + .get_or_init(|| self.connect_once(&self.startup)) + .await + .clone() + .map_err(ExecServerError::ConnectionAttempt) + }; + // Ready may arrive before an older attempt publishes its provisioning failure. + if let Err(ExecServerError::ConnectionAttempt(error)) = &result + && matches!(error.as_ref(), ExecServerError::ProvisioningFailed(_)) + && matches!( + &self.transport_params, + Some(ExecServerTransportParams::Deferred(deferred)) + if matches!(*deferred.readiness.borrow(), Some(Ok(()))) + ) + { + return Box::pin(self.reconnect()).await; + } + result + } + + async fn reconnect(&self) -> Result { + // Callers handling the same outage share one reconnect attempt. + let attempt = { + let mut reconnect = self + .reconnect + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if let Some(client) = self.connected_client() { + return Ok(client); + } + reconnect + .get_or_insert_with(|| Arc::new(ConnectionAttempt::default())) + .clone() + }; + let result = attempt + .result + .get_or_init(|| self.connect_once(&attempt)) + .await; + let mut reconnect = self + .reconnect + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // Forget only this completed attempt so a later operation can retry after failure. + if reconnect + .as_ref() + .is_some_and(|current| Arc::ptr_eq(current, &attempt)) + { + *reconnect = None; + } + result.clone().map_err(ExecServerError::ConnectionAttempt) + } + + fn connected_client(&self) -> Option { + self.cached_client() + .filter(|client| !client.is_disconnected()) + } + + fn cached_client(&self) -> Option { + self.current_client + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .clone() + } + + fn can_reconnect(&self) -> bool { + matches!( + self.transport_params, + Some( + ExecServerTransportParams::Deferred(_) + | ExecServerTransportParams::WebSocketUrl { .. } + | ExecServerTransportParams::NoiseRendezvous { .. } + ) + ) + } +} + +impl HttpClient for LazyRemoteExecServerClient { + fn http_request( + &self, + params: crate::HttpRequestParams, + ) -> BoxFuture<'_, Result> { + async move { self.get().await?.http_request(params).await }.boxed() + } + + fn http_request_stream( + &self, + params: crate::HttpRequestParams, + ) -> BoxFuture< + '_, + Result<(crate::HttpRequestResponse, crate::HttpResponseBodyStream), ExecServerError>, + > { + async move { self.get().await?.http_request_stream(params).await }.boxed() + } +} + +impl LazyRemoteExecServerClient { + pub(crate) async fn environment_info(&self) -> Result { + self.get().await?.environment_info().await + } +} + +#[derive(Debug, thiserror::Error)] +pub enum ExecServerError { + #[error("failed to spawn exec-server: {0}")] + Spawn(#[source] std::io::Error), + #[error("timed out connecting to exec-server websocket `{url}` after {timeout:?}")] + WebSocketConnectTimeout { url: String, timeout: Duration }, + #[error("failed to connect to exec-server websocket `{url}`: {source}")] + WebSocketConnect { + url: String, + #[source] + source: tokio_tungstenite::tungstenite::Error, + }, + #[error("failed to configure exec-server websocket: {0}")] + WebSocketConfiguration(String), + #[error("timed out waiting for exec-server initialize handshake after {timeout:?}")] + InitializeTimedOut { timeout: Duration }, + #[error("exec-server transport closed")] + Closed, + #[error("{0}")] + Disconnected(String), + #[error("environment unavailable: {0}")] + ProvisioningFailed(String), + #[error("failed to serialize or deserialize exec-server JSON: {0}")] + Json(#[from] serde_json::Error), + #[error("HTTP request failed: {0}")] + HttpRequest(String), + #[error("exec-server protocol error: {0}")] + Protocol(String), + #[error( + "environment `{environment_id}` is already registered with a different provisioning mode" + )] + ProvisioningModeConflict { environment_id: String }, + #[error("exec-server rejected request ({code}): {message}")] + Server { code: i64, message: String }, + #[error("environment registry request failed ({status}{code_suffix}): {message}", code_suffix = .code.as_ref().map(|code| format!(", {code}")).unwrap_or_default())] + EnvironmentRegistryHttp { + status: http::StatusCode, + code: Option, + message: String, + }, + #[error("environment registry configuration error: {0}")] + EnvironmentRegistryConfig(String), + #[error("environment registry authentication error: {0}")] + EnvironmentRegistryAuth(String), + #[error("environment registry request failed: {0}")] + EnvironmentRegistryRequest(#[from] codex_http_client::RouteAwareRequestError), + #[error("exec-server connection attempt failed: {0}")] + ConnectionAttempt(#[source] Arc), +} + +impl ExecServerClient { + fn attach_environment_connection_state( + &self, + state_tx: watch::Sender, + ) { + let mut connection = self + .inner + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + connection.environment_connection_state_tx = state_tx; + connection.publish_environment_connection_state(); + } + + fn fail_fast(&self) -> Result { + self.rpc_client_without_recovery()?; + Ok(Self { + inner: Arc::clone(&self.inner), + recovery_policy: RecoveryPolicy::FailFast, + }) + } + + fn rpc_client_without_recovery(&self) -> Result, ExecServerError> { + let connection = self + .inner + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match &connection.status { + ConnectionStatus::Connected(rpc_client) if !rpc_client.is_disconnected() => { + Ok(Arc::clone(rpc_client)) + } + ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => Err( + ExecServerError::Disconnected("exec-server environment is recovering".to_string()), + ), + ConnectionStatus::Failed(message) => { + Err(ExecServerError::Disconnected(message.clone())) + } + } + } + + async fn rpc_client(&self) -> Result, ExecServerError> { + match self.recovery_policy { + RecoveryPolicy::Wait => self.inner.rpc_client().await, + RecoveryPolicy::FailFast => self.rpc_client_without_recovery(), + } + } + + async fn initialize_rpc( + &self, + rpc_client: &RpcClient, + options: ExecServerClientConnectOptions, + noise_context: Option, + ) -> Result { + let ExecServerClientConnectOptions { + client_name, + initialize_timeout, + resume_session_id, + } = options; + let timeout_for_error = noise_context + .as_ref() + .map_or(initialize_timeout, |context| context.timeout_for_error); + + timeout(initialize_timeout, async { + let params = InitializeParams { + client_name, + resume_session_id, + }; + let response: InitializeResponse = if let Some(noise_context) = noise_context { + // This is the one RPC whose wire method and trace operation + // intentionally differ: preserve the compatibility initialize + // parent while measuring only the actual RPC as initialize_rpc. + let initialize_rpc_span = tracing::info_span!( + parent: &noise_context.span, + "codex.exec_server.remote.initialize_rpc", + otel.kind = "client", + otel.name = "codex.exec_server.remote.initialize_rpc", + ); + let response = rpc_client + .call_untraced(INITIALIZE_METHOD, ¶ms) + .instrument(initialize_rpc_span) + .await; + drop(noise_context); + response? + } else { + rpc_client.call(INITIALIZE_METHOD, ¶ms).await? + }; + let session_id = self + .inner + .session_id + .get_or_init(|| response.session_id.clone()); + if session_id != &response.session_id { + return Err(ExecServerError::Protocol(format!( + "exec-server initialized an unexpected session {}", + response.session_id + ))); + } + rpc_client + .notify(INITIALIZED_METHOD, &serde_json::json!({})) + .await?; + Ok(response) + }) + .await + .map_err(|_| ExecServerError::InitializeTimedOut { + timeout: timeout_for_error, + })? + } + + pub async fn exec(&self, params: ExecParams) -> Result { + self.call(EXEC_METHOD, ¶ms).await + } + + /// Returns cached executor metadata, fetching it lazily if initialization omitted it. + pub async fn environment_info(&self) -> Result { + self.inner + .environment_info + .get_or_try_init(|| self.force_environment_info()) + .await + .cloned() + } + + /// Fetches executor metadata over RPC without reading or updating the cache. + // TODO: Remove after app-server migrates off this call. + pub async fn force_environment_info(&self) -> Result { + let rpc_client = self.rpc_client().await?; + self.map_rpc_call_result( + rpc_client + .call_with_timeout(ENVIRONMENT_INFO_METHOD, &(), ENVIRONMENT_INFO_TIMEOUT) + .await, + ) + } + + pub async fn read_environment_config( + &self, + params: EnvironmentConfigReadParams, + ) -> Result { + self.call(ENVIRONMENT_CONFIG_READ_METHOD, ¶ms).await + } + + pub async fn environment_status(&self) -> Result { + // Health checks only reuse an existing RPC connection and never initiate recovery. + let rpc_client = self.rpc_client_without_recovery()?; + self.map_rpc_call_result( + rpc_client + .call_with_timeout(ENVIRONMENT_STATUS_METHOD, &(), ENVIRONMENT_STATUS_TIMEOUT) + .await, + ) + } + + pub async fn discover_capability_roots( + &self, + params: CapabilityRootsDiscoverParams, + ) -> Result { + self.call(CAPABILITY_ROOTS_DISCOVER_METHOD, ¶ms).await + } + + pub async fn read(&self, params: ReadParams) -> Result { + self.call(EXEC_READ_METHOD, ¶ms).await + } + + pub async fn write( + &self, + process_id: &ProcessId, + chunk: Vec, + write_id: String, + ) -> Result { + self.call( + EXEC_WRITE_METHOD, + &WriteParams { + process_id: process_id.clone(), + chunk: chunk.into(), + write_id, + }, + ) + .await + } + + pub async fn signal( + &self, + process_id: &ProcessId, + signal: ProcessSignal, + ) -> Result<(), ExecServerError> { + let _response: SignalResponse = self + .call( + EXEC_SIGNAL_METHOD, + &SignalParams { + process_id: process_id.clone(), + signal, + }, + ) + .await?; + Ok(()) + } + + pub async fn terminate( + &self, + process_id: &ProcessId, + ) -> Result { + // A close notification may arrive before the termination response. + if let Some(session) = self.inner.get_session(process_id) { + session + .network_policy + .cancellation + .record(NetworkRequestCancellationReason::ProcessCancelled); + } + self.call_for_cleanup( + EXEC_TERMINATE_METHOD, + &TerminateParams { + process_id: process_id.clone(), + }, + ) + .await + } + + pub async fn fs_read_file( + &self, + params: FsReadFileParams, + ) -> Result { + self.call(FS_READ_FILE_METHOD, ¶ms).await + } + + pub async fn fs_open(&self, params: FsOpenParams) -> Result { + self.call(FS_OPEN_METHOD, ¶ms).await + } + + pub async fn fs_read_block( + &self, + params: FsReadBlockParams, + ) -> Result { + self.call(FS_READ_BLOCK_METHOD, ¶ms).await + } + + pub async fn fs_close( + &self, + params: FsCloseParams, + ) -> Result { + self.call_for_cleanup(FS_CLOSE_METHOD, ¶ms).await + } + + pub async fn fs_write_file( + &self, + params: FsWriteFileParams, + ) -> Result { + self.call(FS_WRITE_FILE_METHOD, ¶ms).await + } + + pub async fn fs_create_directory( + &self, + params: FsCreateDirectoryParams, + ) -> Result { + self.call(FS_CREATE_DIRECTORY_METHOD, ¶ms).await + } + + pub async fn fs_get_metadata( + &self, + params: FsGetMetadataParams, + ) -> Result { + self.call(FS_GET_METADATA_METHOD, ¶ms).await + } + + pub async fn fs_canonicalize( + &self, + params: FsCanonicalizeParams, + ) -> Result { + self.call(FS_CANONICALIZE_METHOD, ¶ms).await + } + + pub async fn fs_read_directory( + &self, + params: FsReadDirectoryParams, + ) -> Result { + self.call(FS_READ_DIRECTORY_METHOD, ¶ms).await + } + + pub async fn fs_walk(&self, params: FsWalkParams) -> Result { + self.call(FS_WALK_METHOD, ¶ms).await + } + + pub async fn fs_remove( + &self, + params: FsRemoveParams, + ) -> Result { + self.call(FS_REMOVE_METHOD, ¶ms).await + } + + pub async fn fs_copy(&self, params: FsCopyParams) -> Result { + self.call(FS_COPY_METHOD, ¶ms).await + } + + pub(crate) async fn start_process( + &self, + params: ExecParams, + network_policy_decider: Option>, + ) -> Result { + let policy_decision_timeout_ms = params + .network_proxy + .as_ref() + .and_then(|launch| launch.policy_decision_timeout_ms); + let network_policy_controller = match ( + network_policy_decider.as_ref(), + policy_decision_timeout_ms, + ) { + (None, None) => None, + (Some(decider), Some(timeout_ms)) if timeout_ms > 0 => { + Some(NetworkPolicyDecisionController { + decider: Arc::clone(decider), + timeout: Duration::from_millis(timeout_ms), + }) + } + _ => { + return Err(ExecServerError::Protocol( + "network policy decision callback timeout must match the configured decider and be nonzero" + .to_string(), + )); + } + }; + + loop { + let rpc_client = self.rpc_client().await?; + if !self.inner.begin_process_start(&rpc_client) { + continue; + } + + let process_id = params.process_id.clone(); + let mut state = SessionState::new(/*recoverable*/ false); + state.network_policy.audit = + params + .network_proxy + .as_ref() + .map(|launch| NetworkPolicyAuditContext { + metadata: launch.audit_metadata.clone(), + execution_id: launch.execution_id.clone(), + }); + let state = Arc::new(state); + if let Some(controller) = network_policy_controller.as_ref() { + state + .network_policy + .controller + .store(Some(Arc::new(controller.clone()))); + } + if let Err(error) = self.inner.insert_session(&process_id, Arc::clone(&state)) { + self.inner.finish_process_start(); + return Err(error); + } + let active_start = ActiveProcessStart { + inner: Arc::clone(&self.inner), + }; + let mut pending_start = PendingProcessStartSession { + inner: Arc::clone(&self.inner), + process_id: process_id.clone(), + state: Arc::clone(&state), + armed: true, + }; + let client = self.clone(); + let (result_tx, result_rx) = tokio::sync::oneshot::channel(); + let (result_received_tx, result_received_rx) = tokio::sync::oneshot::channel(); + let process_start_task = async move { + let _active_start = active_start; + match client + .call_rpc::<_, ExecResponse>(&rpc_client, EXEC_METHOD, ¶ms) + .await + { + Ok(response) => { + state.recoverable.store(true, Ordering::Release); + let session = Session { + client: client.clone(), + process_id: process_id.clone(), + sandbox_type: response.sandbox_type, + state: Arc::clone(&state), + }; + // Wait for caller receipt so cancellation after send still triggers cleanup. + if result_tx.send(Ok(session)).is_err() || result_received_rx.await.is_err() + { + state.recoverable.store(false, Ordering::Release); + tokio::spawn(async move { + cleanup_process_start(&client, &process_id, &state).await; + }); + } + } + Err(error) => { + if is_transport_closed_error(&error) { + tokio::spawn(async move { + cleanup_process_start(&client, &process_id, &state).await; + }); + } else { + client.inner.remove_session_if(&process_id, &state); + } + let _ = result_tx.send(Err(error)); + } + } + }; + tokio::spawn( + process_start_task + .in_current_span() + .with_current_subscriber(), + ); + let result = result_rx.await; + // The response task may have queued a session before retirement. + if self.inner.retired.is_cancelled() { + return Err(ExecServerError::Disconnected( + "exec-server executor was replaced".to_string(), + )); + } + if matches!(&result, Ok(Ok(_))) { + pending_start.armed = false; + let _ = result_received_tx.send(()); + } + return result.map_err(|_| { + ExecServerError::Protocol("process start task stopped unexpectedly".to_string()) + })?; + } + } + + #[cfg(test)] + pub(crate) async fn register_session( + &self, + process_id: &ProcessId, + ) -> Result { + let state = Arc::new(SessionState::new(/*recoverable*/ true)); + self.inner.insert_session(process_id, Arc::clone(&state))?; + Ok(Session { + client: self.clone(), + process_id: process_id.clone(), + sandbox_type: None, + state, + }) + } + + pub fn session_id(&self) -> Option { + self.inner.session_id.get().cloned() + } + + fn is_disconnected(&self) -> bool { + self.inner.is_failed() + } + + fn readiness_result(&self) -> Option> { + let connection = self + .inner + .connection + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + match &connection.status { + ConnectionStatus::Connected(rpc_client) if !rpc_client.is_disconnected() => { + Some(Ok(())) + } + ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => None, + ConnectionStatus::Failed(message) => { + Some(Err(ExecServerError::Disconnected(message.clone()))) + } + } + } + + pub(crate) async fn connect( + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + ) -> Result { + Self::connect_with_recovery(connection, options, /*reconnect_strategy*/ None).await + } + + pub(crate) async fn connect_with_recovery( + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + reconnect_strategy: Option, + ) -> Result { + Self::connect_with_recovery_inner( + connection, + options, + reconnect_strategy, + /*noise_context*/ None, + ) + .await + } + + pub(crate) async fn connect_with_recovery_and_noise_context( + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + reconnect_strategy: Option, + noise_context: NoiseInitializeContext, + ) -> Result { + Self::connect_with_recovery_inner( + connection, + options, + reconnect_strategy, + Some(noise_context), + ) + .await + } + + async fn connect_with_recovery_inner( + connection: JsonRpcConnection, + options: ExecServerClientConnectOptions, + reconnect_strategy: Option, + noise_context: Option, + ) -> Result { + let (rpc_client, events_rx) = RpcClient::new(connection); + let rpc_client = Arc::new(rpc_client); + let session_id = OnceLock::new(); + let (connection_changed, _connection_changed_rx) = watch::channel(()); + let inner = Arc::new(Inner { + connection: StdMutex::new(ConnectionState { + status: ConnectionStatus::Connected(Arc::clone(&rpc_client)), + active_process_starts: 0, + environment_connection_state_tx: watch::channel( + EnvironmentConnectionState::Connected, + ) + .0, + }), + connection_changed, + sessions: ArcSwap::from_pointee(HashMap::new()), + sessions_write_lock: StdMutex::new(()), + http_body_streams: ArcSwap::from_pointee(HashMap::new()), + http_body_stream_failures: ArcSwap::from_pointee(HashMap::new()), + http_body_streams_write_lock: Mutex::new(()), + http_body_stream_byte_budget: Arc::new(Semaphore::new(MAX_QUEUED_HTTP_BODY_BYTES)), + http_body_stream_next_id: AtomicU64::new(1), + rpc_inbound_request_slots: Arc::new(Semaphore::new(MAX_IN_FLIGHT_SERVER_CALLS)), + session_id, + retired: CancellationToken::new(), + environment_info: OnceCell::new(), + reconnect_strategy, + }); + let client = Self { + inner, + recovery_policy: RecoveryPolicy::Wait, + }; + // An explicit resume can redirect notifications from running processes + // before initialize returns. Drain them immediately so a burst cannot + // fill the bounded event channel and block the initialize response. + client.spawn_rpc_reader(&rpc_client, events_rx); + let initialize_response = client + .initialize_rpc(&rpc_client, options, noise_context) + .await?; + if let Some(info) = initialize_response.environment_info { + assert!( + client.inner.environment_info.set(info).is_ok(), + "new client metadata cache must be empty" + ); + } + Ok(client) + } + + async fn call(&self, method: &str, params: &P) -> Result + where + P: serde::Serialize, + T: serde::de::DeserializeOwned, + { + let rpc_client = self.rpc_client().await?; + self.call_rpc(&rpc_client, method, params).await + } + + async fn call_rpc( + &self, + rpc_client: &Arc, + method: &str, + params: &P, + ) -> Result + where + P: serde::Serialize, + T: serde::de::DeserializeOwned, + { + self.map_rpc_call_result(rpc_client.call(method, params).await) + } + + async fn call_for_cleanup(&self, method: &str, params: &P) -> Result + where + P: serde::Serialize, + T: serde::de::DeserializeOwned, + { + let rpc_client = self.inner.rpc_client().await?; + self.map_rpc_call_result(rpc_client.call_for_cleanup(method, params).await) + } + + fn map_rpc_call_result( + &self, + result: Result, + ) -> Result { + // Explicit retirement rejects late responses. Ordinary EOF still preserves + // responses received before disconnect, as ordered by the RPC reader. + if self.inner.retired.is_cancelled() { + return Err(ExecServerError::Disconnected( + "exec-server executor was replaced".to_string(), + )); + } + result.map_err(|error| { + let error = ExecServerError::from(error); + if is_transport_closed_error(&error) { + ExecServerError::Disconnected(disconnected_message(/*reason*/ None)) + } else { + error + } + }) + } +} + +async fn cleanup_process_start( + client: &ExecServerClient, + process_id: &ProcessId, + state: &Arc, +) { + loop { + match client.terminate(process_id).await { + Ok(_) => break, + Err(error) if is_transport_closed_error(&error) && !client.inner.is_failed() => { + continue; + } + Err(_) => break, + } + } + client.inner.remove_session_if(process_id, state); +} + +impl From for ExecServerError { + fn from(value: RpcCallError) -> Self { + match value { + RpcCallError::Closed => Self::Closed, + RpcCallError::Json(err) => Self::Json(err), + RpcCallError::Server(error) => Self::Server { + code: error.code, + message: error.message, + }, + RpcCallError::TimedOut { method, timeout } => Self::Protocol(format!( + "timed out waiting for exec-server `{method}` response after {timeout:?}" + )), + RpcCallError::PendingRequestLimitExceeded { limit } => Self::Protocol(format!( + "exec-server has reached its limit of {limit} pending requests" + )), + } + } +} + +impl SessionState { + fn new(recoverable: bool) -> Self { + let (wake_tx, _wake_rx) = watch::channel(0); + Self { + wake_tx, + events: ExecProcessEventLog::new( + PROCESS_EVENT_CHANNEL_CAPACITY, + PROCESS_EVENT_RETAINED_BYTES, + ), + ordered_events: StdMutex::new(OrderedSessionEvents::default()), + recoverable: AtomicBool::new(recoverable), + next_write_id: AtomicU64::new(1), + network_policy: NetworkPolicyState { + controller: ArcSwapOption::empty(), + cancelled: CancellationToken::new(), + cancellation: NetworkRequestCancellation::default(), + audit: None, + }, + } + } + + pub(crate) fn subscribe(&self) -> watch::Receiver { + self.wake_tx.subscribe() + } + + pub(crate) fn subscribe_events(&self) -> ExecProcessEventReceiver { + self.events.subscribe() + } + + fn note_change(&self, seq: u64) { + self.wake_tx + .send_modify(|current| *current = (*current).max(seq)); + } + + /// Publishes a process event only when all earlier sequenced events have + /// already been published. + /// + /// Returns `true` only when this call actually publishes the ordered + /// `Closed` event. The caller uses that signal to remove the session route + /// after the terminal event is visible to subscribers, rather than when a + /// possibly-early closed notification first arrives. + fn publish_ordered_event(&self, event: ExecProcessEvent) -> Result { + let Some(seq) = event.seq() else { + self.events.publish(event); + return Ok(false); + }; + + let mut ordered_events = self + .ordered_events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // We have already delivered this sequence number or moved past it, + // so accepting it again would duplicate output or lifecycle events. + if ordered_events.failure.is_some() + || ordered_events.closed_published + || seq <= ordered_events.last_published_seq + { + return Ok(false); + } + + ordered_events.insert_pending(event)?; + Ok(self.publish_ready(&mut ordered_events)) + } + + fn publish_ready(&self, ordered_events: &mut OrderedSessionEvents) -> bool { + let mut published_closed = false; + loop { + let next_seq = ordered_events.last_published_seq.saturating_add(1); + let Some(event) = ordered_events.pending.remove(&next_seq) else { + break; + }; + ordered_events.pending_bytes = ordered_events + .pending_bytes + .saturating_sub(pending_process_event_bytes(&event)); + ordered_events.last_published_seq = next_seq; + ordered_events.exit_published |= matches!(&event, ExecProcessEvent::Exited { .. }); + let is_closed = matches!(&event, ExecProcessEvent::Closed { .. }); + ordered_events.closed_published |= is_closed; + published_closed |= is_closed; + if is_closed && ordered_events.exit_published { + self.network_policy + .cancellation + .record(NetworkRequestCancellationReason::ProcessFinished); + } + self.events.publish(event); + } + published_closed + } + + fn set_failure(&self, message: String) { + let mut ordered_events = self + .ordered_events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + if ordered_events.failure.is_some() || ordered_events.closed_published { + return; + } + ordered_events.failure = Some(message.clone()); + ordered_events.pending.clear(); + ordered_events.pending_bytes = 0; + self.events.publish(ExecProcessEvent::Failed(message)); + drop(ordered_events); + self.wake_tx + .send_modify(|current| *current = current.saturating_add(1)); + } + + fn failed_response(&self) -> Option { + self.ordered_events + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner) + .failure + .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), + sandbox_denied: false, + } + } + + fn next_write_id(&self) -> String { + self.next_write_id + .fetch_add(1, Ordering::Relaxed) + .to_string() + } +} + +impl OrderedSessionEvents { + fn insert_pending(&mut self, event: ExecProcessEvent) -> Result<(), String> { + let Some(seq) = event.seq() else { + return Err("cannot reorder an unsequenced process event".to_string()); + }; + if self.pending.contains_key(&seq) { + return Ok(()); + } + + let next_seq = self.last_published_seq.saturating_add(1); + // The next expected event is synchronously published by every caller, + // so it can drain a full buffer without becoming retained state. + let closes_gap = seq == next_seq; + if !closes_gap && self.pending.len() >= MAX_PENDING_PROCESS_EVENTS { + return Err(format!( + "process event reorder buffer exceeds {MAX_PENDING_PROCESS_EVENTS} entries" + )); + } + + let event_bytes = pending_process_event_bytes(&event); + if event_bytes > MAX_PENDING_PROCESS_EVENT_BYTES { + return Err(format!( + "process event exceeds {MAX_PENDING_PROCESS_EVENT_BYTES} bytes" + )); + } + let pending_bytes = self.pending_bytes.saturating_add(event_bytes); + if !closes_gap && pending_bytes > MAX_PENDING_PROCESS_EVENT_BYTES { + return Err(format!( + "process event reorder buffer exceeds {MAX_PENDING_PROCESS_EVENT_BYTES} bytes" + )); + } + + self.pending.insert(seq, event); + self.pending_bytes = pending_bytes; + Ok(()) + } +} + +fn pending_process_event_bytes(event: &ExecProcessEvent) -> usize { + match event { + ExecProcessEvent::Output(chunk) => chunk.chunk.0.len(), + ExecProcessEvent::Failed(message) => message.len(), + ExecProcessEvent::Exited { .. } | ExecProcessEvent::Closed { .. } => 0, + } +} + +fn finish_process_event( + inner: &Inner, + process_id: &ProcessId, + session: &Arc, + result: Result, +) { + match result { + Ok(true) => inner.remove_session_if(process_id, session), + Ok(false) => {} + Err(message) => { + session.set_failure(message); + inner.remove_session_if(process_id, session); + } + } +} + +impl Session { + pub(crate) fn process_id(&self) -> &ProcessId { + &self.process_id + } + + pub(crate) fn sandbox_type(&self) -> Option { + self.sandbox_type + } + + pub(crate) fn subscribe_wake(&self) -> watch::Receiver { + self.state.subscribe() + } + + pub(crate) fn subscribe_events(&self) -> ExecProcessEventReceiver { + self.state.subscribe_events() + } + + pub(crate) async fn read( + &self, + after_seq: Option, + max_bytes: Option, + wait_ms: Option, + ) -> Result { + loop { + if let Some(response) = self.state.failed_response() { + return Ok(response); + } + + match self + .client + .read(ReadParams { + process_id: self.process_id.clone(), + after_seq, + max_bytes, + wait_ms, + }) + .await + { + Ok(response) => return Ok(response), + Err(error) + if is_transport_closed_error(&error) && !self.client.inner.is_failed() => + { + continue; + } + Err(error) if is_transport_closed_error(&error) => { + if let Some(response) = self.state.failed_response() { + return Ok(response); + } + let message = error.to_string(); + self.state.set_failure(message.clone()); + return Ok(self.state.synthesized_failure(message)); + } + Err(error) => return Err(error), + } + } + } + + pub(crate) async fn write(&self, chunk: Vec) -> Result { + let write_id = self.state.next_write_id(); + loop { + match self + .client + .write(&self.process_id, chunk.clone(), write_id.clone()) + .await + { + Ok(response) => return Ok(response), + Err(error) + if is_transport_closed_error(&error) && !self.client.inner.is_failed() => + { + continue; + } + Err(error) => return Err(error), + } + } + } + + pub(crate) async fn signal(&self, signal: ProcessSignal) -> Result<(), ExecServerError> { + self.client.signal(&self.process_id, signal).await + } + + pub(crate) async fn terminate(&self) -> Result<(), ExecServerError> { + self.client.terminate(&self.process_id).await?; + self.cancel_network_policy_decisions(); + Ok(()) + } + + pub(crate) fn cancel_network_policy_decisions(&self) { + self.state + .network_policy + .cancellation + .record(NetworkRequestCancellationReason::ProcessCancelled); + self.state.network_policy.cancelled.cancel(); + } + + pub(crate) async fn unregister(&self) { + self.client + .inner + .remove_session_if(&self.process_id, &self.state); + } +} + +impl Inner { + fn get_session(&self, process_id: &ProcessId) -> Option> { + self.sessions.load().get(process_id).cloned() + } + + fn insert_session( + &self, + process_id: &ProcessId, + session: Arc, + ) -> Result<(), ExecServerError> { + let _sessions_write_guard = self + .sessions_write_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + // Do not register a process session that can never receive environment + // notifications. Without this check, remote MCP startup could create a + // dead session and wait for process output that will never arrive. + if let Some(message) = self.failure_message() { + return Err(ExecServerError::Disconnected(message)); + } + 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(()) + } + + fn remove_session_if(&self, process_id: &ProcessId, expected: &Arc) { + let _sessions_write_guard = self + .sessions_write_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + let sessions = self.sessions.load(); + if !sessions + .get(process_id) + .is_some_and(|session| Arc::ptr_eq(session, expected)) + { + return; + } + let mut next_sessions = sessions.as_ref().clone(); + next_sessions.remove(process_id); + self.sessions.store(Arc::new(next_sessions)); + expected + .network_policy + .cancellation + .record(NetworkRequestCancellationReason::ProcessCancelled); + expected.network_policy.cancelled.cancel(); + expected.network_policy.controller.store(None); + } + + fn take_all_sessions(&self) -> HashMap> { + let _sessions_write_guard = self + .sessions_write_lock + .lock() + .unwrap_or_else(std::sync::PoisonError::into_inner); + 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 | ExecServerError::Disconnected(_) + ) || matches!( + error, + ExecServerError::Server { + code: -32000, + message, + } if message == "JSON-RPC transport closed" + ) +} + +fn fail_all_sessions(inner: &Arc, message: String) { + let sessions = inner.take_all_sessions(); + + for (_, session) in sessions { + session + .network_policy + .cancellation + .record(NetworkRequestCancellationReason::ConnectionClosed); + session.network_policy.cancelled.cancel(); + session.network_policy.controller.store(None); + // Sessions synthesize a closed read response and emit a pushed Failed + // event. That covers both polling consumers and streaming consumers + // such as environment-backed MCP stdio. + session.set_failure(message.clone()); + } +} + +/// Fails all in-flight work that depends on the shared JSON-RPC transport. +async fn fail_all_in_flight_work(inner: &Arc, message: String) { + fail_all_sessions(inner, message.clone()); + inner.fail_all_http_body_streams(message).await; +} + +async fn handle_server_notification( + inner: &Arc, + notification: JSONRPCNotification, +) -> Result<(), ExecServerError> { + match notification.method.as_str() { + EXEC_OUTPUT_DELTA_METHOD => { + let params: ExecOutputDeltaNotification = + serde_json::from_value(notification.params.unwrap_or(Value::Null))?; + if let Some(session) = inner.get_session(¶ms.process_id) { + let result = + session.publish_ordered_event(ExecProcessEvent::Output(ProcessOutputChunk { + seq: params.seq, + stream: params.stream, + chunk: params.chunk, + })); + if result.is_ok() { + session.note_change(params.seq); + } + finish_process_event(inner, ¶ms.process_id, &session, result); + } + } + EXEC_EXITED_METHOD => { + let params: ExecExitedNotification = + serde_json::from_value(notification.params.unwrap_or(Value::Null))?; + if let Some(session) = inner.get_session(¶ms.process_id) { + let result = session.publish_ordered_event(ExecProcessEvent::Exited { + seq: params.seq, + exit_code: params.exit_code, + sandbox_denied: params.sandbox_denied, + }); + if result.is_ok() { + session.note_change(params.seq); + } + finish_process_event(inner, ¶ms.process_id, &session, result); + } + } + EXEC_CLOSED_METHOD => { + let params: ExecClosedNotification = + serde_json::from_value(notification.params.unwrap_or(Value::Null))?; + if let Some(session) = inner.get_session(¶ms.process_id) { + // Closed is terminal, but it can arrive before tail output or + // exited. Keep routing this process until the ordered publisher + // says Closed has actually been delivered. + let result = + session.publish_ordered_event(ExecProcessEvent::Closed { seq: params.seq }); + if result.is_ok() { + session.note_change(params.seq); + } + finish_process_event(inner, ¶ms.process_id, &session, result); + } + } + HTTP_REQUEST_BODY_DELTA_METHOD => { + inner + .handle_http_body_delta_notification(notification.params) + .await?; + } + NETWORK_POLICY_DECISION_METHOD => { + let Ok(params) = serde_json::from_value::( + notification.params.unwrap_or(Value::Null), + ) else { + debug!("ignoring malformed exec-server network policy decision notification"); + return Ok(()); + }; + let Some(session) = inner.get_session(¶ms.process_id) else { + debug!("ignoring network policy decision for an unknown exec-server process"); + return Ok(()); + }; + let Some(context) = session.network_policy.audit.as_ref() else { + return Ok(()); + }; + if !network_policy_audit::emit_network_policy_decision(context, ¶ms) { + debug!("ignoring invalid exec-server network policy decision notification"); + } + } + other => { + debug!("ignoring unknown exec-server notification: {other}"); + } + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use codex_exec_server_protocol::JSONRPCMessage; + use codex_exec_server_protocol::JSONRPCNotification; + use codex_exec_server_protocol::JSONRPCResponse; + use codex_http_client::HttpClientFactory; + use codex_http_client::OutboundProxyPolicy; + use codex_utils_path_uri::PathUri; + use futures::SinkExt; + use futures::StreamExt; + use http::HeaderMap; + use opentelemetry::trace::TracerProvider as _; + use opentelemetry_sdk::trace::SdkTracerProvider; + use pretty_assertions::assert_eq; + use std::collections::HashMap; + #[cfg(unix)] + use std::path::Path; + #[cfg(unix)] + use std::process::Command; + use std::sync::Arc; + use tokio::io::AsyncBufReadExt; + use tokio::io::AsyncWrite; + use tokio::io::AsyncWriteExt; + use tokio::io::BufReader; + use tokio::io::duplex; + use tokio::net::TcpListener; + use tokio::net::TcpStream; + use tokio::sync::mpsc; + use tokio::sync::oneshot; + use tokio::sync::watch; + use tokio::time::Duration; + #[cfg(unix)] + use tokio::time::sleep; + use tokio::time::timeout; + use tokio_tungstenite::WebSocketStream; + use tokio_tungstenite::accept_async; + use tokio_tungstenite::tungstenite::Message; + use tracing::Instrument; + use tracing_subscriber::filter::filter_fn; + use tracing_subscriber::prelude::*; + + 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; + use crate::client_api::ExecServerTransportParams; + use crate::client_api::RemoteExecServerConnectArgs; + use crate::client_api::StdioExecServerCommand; + use crate::client_api::StdioExecServerConnectArgs; + use crate::connection::JsonRpcConnection; + use crate::process::ExecProcessEvent; + 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_WRITE_METHOD; + use crate::protocol::EnvironmentInfo; + use crate::protocol::ExecClosedNotification; + use crate::protocol::ExecExitedNotification; + use crate::protocol::ExecOutputDeltaNotification; + use crate::protocol::ExecOutputStream; + use crate::protocol::ExecParams; + use crate::protocol::ExecResponse; + use crate::protocol::INITIALIZE_METHOD; + use crate::protocol::INITIALIZED_METHOD; + use crate::protocol::InitializeResponse; + use crate::protocol::ProcessOutputChunk; + use crate::protocol::ProcessSandboxType; + use crate::protocol::ReadResponse; + use crate::protocol::WriteParams; + use crate::protocol::WriteResponse; + use crate::protocol::WriteStatus; + + async fn read_jsonrpc_line(lines: &mut tokio::io::Lines>) -> 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(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(flavor = "multi_thread", worker_threads = 1)] + async fn process_start_propagates_caller_trace_context_across_background_task() { + let (client_stdin, server_reader) = duplex(1 << 20); + let (mut server_writer, client_stdout) = duplex(1 << 20); + let server = tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + let initialize = read_jsonrpc_line(&mut lines).await; + let initialize = 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: initialize.id, + result: serde_json::to_value(InitializeResponse { + session_id: "trace-test".to_string(), + environment_info: None, + }) + .expect("initialize response should serialize"), + }), + ) + .await; + + match read_jsonrpc_line(&mut lines).await { + JSONRPCMessage::Notification(notification) + if notification.method == INITIALIZED_METHOD => {} + other => panic!("expected initialized notification, got {other:?}"), + } + + let request = match read_jsonrpc_line(&mut lines).await { + JSONRPCMessage::Request(request) if request.method == EXEC_METHOD => request, + other => panic!("expected process start request, got {other:?}"), + }; + let trace = request.trace.clone(); + let params: ExecParams = + serde_json::from_value(request.params.expect("process start params should exist")) + .expect("process start params should deserialize"); + write_jsonrpc_line( + &mut server_writer, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(ExecResponse { + process_id: params.process_id, + sandbox_type: Some(ProcessSandboxType::LinuxSeccomp), + }) + .expect("process start response should serialize"), + }), + ) + .await; + trace + }); + + let client = ExecServerClient::connect( + JsonRpcConnection::from_stdio( + client_stdout, + client_stdin, + "trace-test-client".to_string(), + ), + ExecServerClientConnectOptions::default(), + ) + .await + .expect("client should connect"); + + let tracer_provider = SdkTracerProvider::builder().build(); + let tracer = tracer_provider.tracer("exec-server-test"); + let subscriber = tracing_subscriber::registry().with( + tracing_opentelemetry::layer() + .with_tracer(tracer) + .with_filter(filter_fn(codex_otel::OtelProvider::trace_export_filter)), + ); + let _subscriber_guard = tracing::subscriber::set_default(subscriber); + tracing::callsite::rebuild_interest_cache(); + let parent_span = tracing::info_span!("process-start-parent"); + let expected_trace = codex_otel::span_w3c_trace_context(&parent_span) + .expect("parent span should have trace context"); + let process_id = ProcessId::from("trace-process"); + + let session = client + .start_process( + ExecParams { + metadata: Default::default(), + process_id: process_id.clone(), + argv: vec!["true".to_string()], + cwd: PathUri::from_host_native_path(std::env::current_dir().expect("cwd")) + .expect("cwd URI"), + shell_snapshot: None, + env_policy: None, + env: HashMap::new(), + tty: false, + pipe_stdin: false, + arg0: None, + sandbox: None, + enforce_managed_network: false, + managed_network: None, + network_proxy: None, + }, + /*network_policy_decider*/ None, + ) + .instrument(parent_span) + .await + .expect("process start should succeed"); + + assert_eq!(session.process_id(), &process_id); + assert_eq!( + session.sandbox_type(), + Some(ProcessSandboxType::LinuxSeccomp) + ); + let trace = server.await.expect("server task").expect("trace context"); + let expected_traceparent = expected_trace + .traceparent + .as_deref() + .expect("parent traceparent"); + let traceparent = trace.traceparent.as_deref().expect("request traceparent"); + let expected_parts = expected_traceparent.split('-').collect::>(); + let parts = traceparent.split('-').collect::>(); + assert_eq!(parts[1], expected_parts[1]); + assert_ne!(parts[2], expected_parts[2]); + assert_eq!(trace.tracestate, expected_trace.tracestate); + } + + async fn accept_websocket(listener: &TcpListener) -> WebSocketStream { + let (stream, _) = listener.accept().await.expect("listener should accept"); + accept_async(stream) + .await + .expect("websocket handshake should succeed") + } + + async fn read_jsonrpc_websocket(websocket: &mut WebSocketStream) -> JSONRPCMessage { + loop { + match timeout(Duration::from_secs(1), websocket.next()) + .await + .expect("json-rpc websocket read should not time out") + .expect("websocket should stay open") + .expect("websocket frame should read") + { + Message::Text(text) => { + return serde_json::from_str(text.as_ref()) + .expect("json-rpc text frame should parse"); + } + Message::Binary(bytes) => { + return serde_json::from_slice(bytes.as_ref()) + .expect("json-rpc binary frame should parse"); + } + Message::Ping(_) | Message::Pong(_) => {} + other => panic!("expected json-rpc websocket frame, got {other:?}"), + } + } + } + + async fn write_jsonrpc_websocket( + websocket: &mut WebSocketStream, + message: JSONRPCMessage, + ) { + let encoded = serde_json::to_string(&message).expect("json-rpc should serialize"); + websocket + .send(Message::Text(encoded.into())) + .await + .expect("json-rpc websocket frame should write"); + } + + async fn complete_websocket_initialize( + websocket: &mut WebSocketStream, + session_id: &str, + expected_resume_session_id: Option<&str>, + ) { + complete_websocket_initialize_with_environment_info( + websocket, + session_id, + expected_resume_session_id, + /*environment_info*/ None, + ) + .await; + } + + async fn complete_websocket_initialize_with_environment_info( + websocket: &mut WebSocketStream, + session_id: &str, + expected_resume_session_id: Option<&str>, + environment_info: Option, + ) { + let initialize = read_jsonrpc_websocket(websocket).await; + let request = match initialize { + JSONRPCMessage::Request(request) if request.method == INITIALIZE_METHOD => request, + other => panic!("expected initialize request, got {other:?}"), + }; + let params: crate::protocol::InitializeParams = + serde_json::from_value(request.params.expect("initialize params should exist")) + .expect("initialize params should deserialize"); + assert_eq!( + params.resume_session_id.as_deref(), + expected_resume_session_id + ); + write_jsonrpc_websocket( + websocket, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(InitializeResponse { + session_id: session_id.to_string(), + environment_info, + }) + .expect("initialize response should serialize"), + }), + ) + .await; + + let initialized = read_jsonrpc_websocket(websocket).await; + match initialized { + JSONRPCMessage::Notification(notification) + if notification.method == INITIALIZED_METHOD => {} + other => panic!("expected initialized notification, got {other:?}"), + } + } + + #[cfg(not(windows))] + #[tokio::test] + async fn connect_stdio_command_initializes_json_rpc_client() { + let client = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + command: StdioExecServerCommand { + program: "sh".to_string(), + args: vec![ + "-c".to_string(), + "read _line; printf '%s\\n' '{\"id\":1,\"result\":{\"sessionId\":\"stdio-test\"}}'; read _line; sleep 60".to_string(), + ], + env: HashMap::new(), + cwd: None, + }, + client_name: "stdio-test-client".to_string(), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + }) + .await + .expect("stdio client should connect"); + + assert_eq!(client.session_id().as_deref(), Some("stdio-test")); + } + + #[cfg(not(windows))] + #[tokio::test] + async fn connect_for_transport_initializes_stdio_command() { + let client = ExecServerClient::connect_for_transport( + ExecServerTransportParams::StdioCommand { + command: StdioExecServerCommand { + program: "sh".to_string(), + args: vec![ + "-c".to_string(), + "read _line; printf '%s\\n' '{\"id\":1,\"result\":{\"sessionId\":\"stdio-test\"}}'; read _line; sleep 60".to_string(), + ], + env: HashMap::new(), + cwd: None, + }, + initialize_timeout: DEFAULT_REMOTE_EXEC_SERVER_INITIALIZE_TIMEOUT, + }, + codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + ) + .await + .expect("stdio transport should connect"); + + assert_eq!(client.session_id().as_deref(), Some("stdio-test")); + } + + #[cfg(windows)] + #[tokio::test] + async fn connect_stdio_command_initializes_json_rpc_client_on_windows() { + let client = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + command: StdioExecServerCommand { + program: "powershell".to_string(), + args: vec![ + "-NoProfile".to_string(), + "-Command".to_string(), + "$null = [Console]::In.ReadLine(); [Console]::Out.WriteLine('{\"id\":1,\"result\":{\"sessionId\":\"stdio-test\"}}'); $null = [Console]::In.ReadLine(); Start-Sleep -Seconds 60".to_string(), + ], + env: HashMap::new(), + cwd: None, + }, + client_name: "stdio-test-client".to_string(), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + }) + .await + .expect("stdio client should connect"); + + assert_eq!(client.session_id().as_deref(), Some("stdio-test")); + } + + #[cfg(unix)] + #[tokio::test] + async fn dropping_stdio_client_terminates_spawned_process() { + let tempdir = tempfile::tempdir().expect("tempdir should be created"); + let pid_file = tempdir.path().join("server.pid"); + let child_pid_file = tempdir.path().join("server-child.pid"); + let stdio_script = format!( + "read _line; \ + echo \"$$\" > {}; \ + sleep 60 >/dev/null 2>&1 & echo \"$!\" > {}; \ + printf '%s\\n' '{{\"id\":1,\"result\":{{\"sessionId\":\"stdio-test\"}}}}'; \ + read _line; \ + wait", + shell_quote(pid_file.as_path()), + shell_quote(child_pid_file.as_path()), + ); + + let client = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + command: StdioExecServerCommand { + program: "sh".to_string(), + args: vec!["-c".to_string(), stdio_script], + env: HashMap::new(), + cwd: None, + }, + client_name: "stdio-test-client".to_string(), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + }) + .await + .expect("stdio client should connect"); + let server_pid = read_pid_file(pid_file.as_path()).await; + let child_pid = read_pid_file(child_pid_file.as_path()).await; + assert!( + process_exists(server_pid), + "spawned stdio process should be running before client drop" + ); + assert!( + process_exists(child_pid), + "spawned stdio child process should be running before client drop" + ); + + drop(client); + + wait_for_process_exit(server_pid).await; + wait_for_process_exit(child_pid).await; + } + + #[cfg(unix)] + #[tokio::test] + async fn malformed_stdio_message_terminates_spawned_process() { + let tempdir = tempfile::tempdir().expect("tempdir should be created"); + let pid_file = tempdir.path().join("server.pid"); + let stdio_script = format!( + "read _line; \ + echo \"$$\" > {}; \ + printf '%s\\n' 'not-json'; \ + sleep 60", + shell_quote(pid_file.as_path()), + ); + + let result = ExecServerClient::connect_stdio_command(StdioExecServerConnectArgs { + command: StdioExecServerCommand { + program: "sh".to_string(), + args: vec!["-c".to_string(), stdio_script], + env: HashMap::new(), + cwd: None, + }, + client_name: "stdio-test-client".to_string(), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + }) + .await; + assert!(result.is_err(), "malformed stdio server should not connect"); + + let server_pid = read_pid_file(pid_file.as_path()).await; + wait_for_process_exit(server_pid).await; + } + + #[cfg(unix)] + async fn read_pid_file(path: &Path) -> u32 { + for _ in 0..20 { + if let Ok(contents) = std::fs::read_to_string(path) { + return contents + .trim() + .parse() + .expect("pid file should contain a pid"); + } + sleep(Duration::from_millis(50)).await; + } + panic!("pid file {} should be written", path.display()); + } + + #[cfg(unix)] + async fn wait_for_process_exit(pid: u32) { + for _ in 0..20 { + if !process_exists(pid) { + return; + } + sleep(Duration::from_millis(100)).await; + } + panic!("process {pid} should exit"); + } + + #[cfg(unix)] + fn process_exists(pid: u32) -> bool { + Command::new("kill") + .arg("-0") + .arg(pid.to_string()) + .status() + .is_ok_and(|status| status.success()) + } + + #[cfg(unix)] + fn shell_quote(path: &Path) -> String { + let value = path.to_string_lossy(); + format!("'{}'", value.replace('\'', "'\\''")) + } + + #[tokio::test] + async fn process_events_are_delivered_in_seq_order_when_notifications_are_reordered() { + 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 { + session_id: "session-1".to_string(), + environment_info: None, + }) + .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 process_id = ProcessId::from("reordered"); + let session = client + .register_session(&process_id) + .await + .expect("session should register"); + let mut events = session.subscribe_events(); + + for message in [ + JSONRPCMessage::Notification(JSONRPCNotification { + method: EXEC_CLOSED_METHOD.to_string(), + params: Some( + serde_json::to_value(ExecClosedNotification { + process_id: process_id.clone(), + seq: 4, + }) + .expect("closed notification should serialize"), + ), + }), + JSONRPCMessage::Notification(JSONRPCNotification { + method: EXEC_OUTPUT_DELTA_METHOD.to_string(), + params: Some( + serde_json::to_value(ExecOutputDeltaNotification { + process_id: process_id.clone(), + seq: 1, + stream: ExecOutputStream::Stdout, + chunk: b"one".to_vec().into(), + }) + .expect("output notification should serialize"), + ), + }), + JSONRPCMessage::Notification(JSONRPCNotification { + method: EXEC_EXITED_METHOD.to_string(), + params: Some( + serde_json::to_value(ExecExitedNotification { + process_id: process_id.clone(), + seq: 3, + exit_code: 0, + sandbox_denied: Some(true), + }) + .expect("exit notification should serialize"), + ), + }), + JSONRPCMessage::Notification(JSONRPCNotification { + method: EXEC_OUTPUT_DELTA_METHOD.to_string(), + params: Some( + serde_json::to_value(ExecOutputDeltaNotification { + process_id: process_id.clone(), + seq: 2, + stream: ExecOutputStream::Stderr, + chunk: b"two".to_vec().into(), + }) + .expect("output notification should serialize"), + ), + }), + ] { + notifications_tx + .send(message) + .await + .expect("notification should queue"); + } + + let mut delivered = Vec::new(); + for _ in 0..4 { + delivered.push( + timeout(Duration::from_secs(1), events.recv()) + .await + .expect("process event should not time out") + .expect("process event stream should stay open"), + ); + } + + assert_eq!( + delivered, + vec![ + ExecProcessEvent::Output(ProcessOutputChunk { + seq: 1, + stream: ExecOutputStream::Stdout, + chunk: b"one".to_vec().into(), + }), + ExecProcessEvent::Output(ProcessOutputChunk { + seq: 2, + stream: ExecOutputStream::Stderr, + chunk: b"two".to_vec().into(), + }), + ExecProcessEvent::Exited { + seq: 3, + exit_code: 0, + sandbox_denied: Some(true), + }, + ExecProcessEvent::Closed { seq: 4 }, + ] + ); + + drop(notifications_tx); + drop(client); + server.await.expect("server task should finish"); + } + + #[tokio::test] + async fn transport_disconnect_fails_sessions_and_rejects_new_sessions() { + let (client_stdin, server_reader) = duplex(1 << 20); + let (mut server_writer, client_stdout) = duplex(1 << 20); + let (disconnect_tx, disconnect_rx) = oneshot::channel(); + 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 { + session_id: "session-1".to_string(), + environment_info: None, + }) + .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:?}"), + } + + let _ = disconnect_rx.await; + drop(server_writer); + }); + + 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 process_id = ProcessId::from("disconnect"); + let session = client + .register_session(&process_id) + .await + .expect("session should register"); + let mut events = session.subscribe_events(); + + disconnect_tx.send(()).expect("disconnect should signal"); + + let event = timeout(Duration::from_secs(1), events.recv()) + .await + .expect("session failure should not time out") + .expect("session event stream should stay open"); + let ExecProcessEvent::Failed(message) = event else { + panic!("expected session failure after disconnect, got {event:?}"); + }; + assert_eq!(message, "exec-server transport disconnected"); + + let response = session + .read( + /*after_seq*/ None, /*max_bytes*/ None, /*wait_ms*/ None, + ) + .await + .expect("disconnected session read should synthesize a response"); + assert_eq!( + response.failure.as_deref(), + Some("exec-server transport disconnected") + ); + assert!(response.closed); + + let new_session = client.register_session(&ProcessId::from("new")).await; + assert!(matches!( + new_session, + Err(super::ExecServerError::Disconnected(_)) + )); + + drop(client); + server.await.expect("server task should finish"); + } + + #[test_case::test_case(Some(EnvironmentInfo::local()); "from_initialize")] + #[test_case::test_case(Some(EnvironmentInfo { + executor_version: "1.2.3-alpha.4".to_string(), + provider_id: Some("sha256:fb4f62da3e84f6864dcec8ede7bc66f1c96ecaeaf55f8a786b85df994057c8ac".to_string()), + ..EnvironmentInfo::local() + }); "with_executor_metadata")] + #[test_case::test_case(None; "legacy_server")] + #[tokio::test] + async fn environment_info_is_cached( + initial_environment_info: Option, + ) -> anyhow::Result<()> { + let listener = TcpListener::bind("127.0.0.1:0").await?; + let websocket_url = format!("ws://{}", listener.local_addr()?); + let expected_info = initial_environment_info + .clone() + .unwrap_or_else(EnvironmentInfo::local); + let server_info = expected_info.clone(); + let server = tokio::spawn(async move { + let mut websocket = accept_websocket(&listener).await; + complete_websocket_initialize_with_environment_info( + &mut websocket, + "session-1", + /*expected_resume_session_id*/ None, + initial_environment_info.clone(), + ) + .await; + if initial_environment_info.is_none() { + let JSONRPCMessage::Request(request) = read_jsonrpc_websocket(&mut websocket).await + else { + panic!("expected environment info request"); + }; + assert_eq!(request.method, "environment/info"); + write_jsonrpc_websocket( + &mut websocket, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(server_info) + .expect("environment info should serialize"), + }), + ) + .await; + } + }); + let client = ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { + websocket_url, + client_name: "metadata-test".to_string(), + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + resume_session_id: None, + http_client_factory: HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + }) + .await?; + + assert_eq!(client.environment_info().await?, expected_info); + server.await?; + // The server is gone, so a cloned client must use the shared cache. + assert_eq!(client.clone().environment_info().await?, expected_info); + Ok(()) + } + + #[tokio::test] + async fn remote_websocket_client_resumes_session() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let (resumed_tx, resumed_rx) = oneshot::channel(); + let (finish_tx, finish_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let mut first = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut first, + "session-1", + /*expected_resume_session_id*/ None, + ) + .await; + first.close(None).await.expect("websocket should close"); + + let mut resumed = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut resumed, + "session-1", + /*expected_resume_session_id*/ Some("session-1"), + ) + .await; + resumed_tx.send(()).expect("resume should signal"); + finish_rx.await.expect("test should finish"); + }); + + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + http_headers: HeaderMap::new(), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + let stable_client = client.get().await.expect("client should connect"); + timeout(Duration::from_secs(1), resumed_rx) + .await + .expect("session resume should not time out") + .expect("session resume should signal"); + let reused_client = client.get().await.expect("client should stay connected"); + assert_eq!(stable_client.session_id().as_deref(), Some("session-1")); + assert!(Arc::ptr_eq(&stable_client.inner, &reused_client.inner)); + finish_tx.send(()).expect("test should finish"); + server.await.expect("server task should finish"); + } + + #[tokio::test] + async fn session_write_retries_same_write_id_after_recovery() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let (finish_tx, finish_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let mut first = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut first, + "session-1", + /*expected_resume_session_id*/ None, + ) + .await; + + let first_write = read_jsonrpc_websocket(&mut first).await; + let first_write = match first_write { + JSONRPCMessage::Request(request) if request.method == EXEC_WRITE_METHOD => request, + other => panic!("expected first process/write request, got {other:?}"), + }; + let first_write_params: WriteParams = + serde_json::from_value(first_write.params.expect("write params should exist")) + .expect("write params should deserialize"); + assert_eq!(first_write_params.process_id.as_str(), "proc-write"); + assert_eq!(first_write_params.chunk.into_inner(), b"hello\n".to_vec()); + let write_id = first_write_params.write_id; + assert!(!write_id.is_empty()); + drop(first); + + let mut resumed = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut resumed, + "session-1", + /*expected_resume_session_id*/ Some("session-1"), + ) + .await; + + let recovery_read = read_jsonrpc_websocket(&mut resumed).await; + let recovery_read = match recovery_read { + JSONRPCMessage::Request(request) if request.method == EXEC_READ_METHOD => request, + other => panic!("expected recovery process/read request, got {other:?}"), + }; + write_jsonrpc_websocket( + &mut resumed, + JSONRPCMessage::Response(JSONRPCResponse { + id: recovery_read.id, + result: serde_json::to_value(ReadResponse { + chunks: Vec::new(), + next_seq: 1, + exited: false, + exit_code: None, + closed: false, + failure: None, + sandbox_denied: false, + }) + .expect("read response should serialize"), + }), + ) + .await; + + let retried_write = read_jsonrpc_websocket(&mut resumed).await; + let retried_write = match retried_write { + JSONRPCMessage::Request(request) if request.method == EXEC_WRITE_METHOD => request, + other => panic!("expected retried process/write request, got {other:?}"), + }; + let retried_write_params: WriteParams = + serde_json::from_value(retried_write.params.expect("write params should exist")) + .expect("write params should deserialize"); + assert_eq!(retried_write_params.process_id.as_str(), "proc-write"); + assert_eq!(retried_write_params.chunk.into_inner(), b"hello\n".to_vec()); + assert_eq!(retried_write_params.write_id, write_id); + write_jsonrpc_websocket( + &mut resumed, + JSONRPCMessage::Response(JSONRPCResponse { + id: retried_write.id, + result: serde_json::to_value(WriteResponse { + status: WriteStatus::Accepted, + }) + .expect("write response should serialize"), + }), + ) + .await; + + finish_rx.await.expect("test should finish"); + }); + + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + http_headers: HeaderMap::new(), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + let stable_client = client.get().await.expect("client should connect"); + let session = stable_client + .register_session(&ProcessId::from("proc-write")) + .await + .expect("session should register"); + + let response = timeout(Duration::from_secs(2), session.write(b"hello\n".to_vec())) + .await + .expect("write should not time out") + .expect("write should recover"); + assert_eq!( + response, + WriteResponse { + status: WriteStatus::Accepted + } + ); + + finish_tx.send(()).expect("test should finish"); + server.await.expect("server task should finish"); + } + + #[tokio::test] + async fn explicit_resume_drains_notifications_before_initialize_response() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let (initialized_tx, initialized_rx) = oneshot::channel(); + let (finish_tx, finish_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let mut websocket = accept_websocket(&listener).await; + let initialize = read_jsonrpc_websocket(&mut websocket).await; + let request = match initialize { + JSONRPCMessage::Request(request) if request.method == INITIALIZE_METHOD => request, + other => panic!("expected initialize request, got {other:?}"), + }; + let params: crate::protocol::InitializeParams = + serde_json::from_value(request.params.expect("initialize params should exist")) + .expect("initialize params should deserialize"); + assert_eq!(params.resume_session_id.as_deref(), Some("session-1")); + + for seq in 1..=256 { + write_jsonrpc_websocket( + &mut websocket, + JSONRPCMessage::Notification(JSONRPCNotification { + method: EXEC_OUTPUT_DELTA_METHOD.to_string(), + params: Some( + serde_json::to_value(ExecOutputDeltaNotification { + process_id: ProcessId::from("busy-process"), + seq, + stream: ExecOutputStream::Stdout, + chunk: b"output".to_vec().into(), + }) + .expect("output notification should serialize"), + ), + }), + ) + .await; + } + write_jsonrpc_websocket( + &mut websocket, + JSONRPCMessage::Response(JSONRPCResponse { + id: request.id, + result: serde_json::to_value(InitializeResponse { + session_id: "session-1".to_string(), + environment_info: None, + }) + .expect("initialize response should serialize"), + }), + ) + .await; + + let initialized = read_jsonrpc_websocket(&mut websocket).await; + match initialized { + JSONRPCMessage::Notification(notification) + if notification.method == INITIALIZED_METHOD => {} + other => panic!("expected initialized notification, got {other:?}"), + } + initialized_tx + .send(()) + .expect("initialized notification should signal"); + finish_rx.await.expect("test should finish"); + }); + + let client = timeout( + Duration::from_secs(1), + ExecServerClient::connect_websocket(RemoteExecServerConnectArgs { + websocket_url, + client_name: "test-client".to_string(), + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + resume_session_id: Some("session-1".to_string()), + http_client_factory: codex_http_client::HttpClientFactory::new( + codex_http_client::OutboundProxyPolicy::ReqwestDefault, + ), + }), + ) + .await + .expect("explicit resume should not time out") + .expect("explicit resume should connect"); + assert_eq!(client.session_id().as_deref(), Some("session-1")); + + timeout(Duration::from_secs(1), initialized_rx) + .await + .expect("initialized notification should not time out") + .expect("initialized notification should signal"); + finish_tx.send(()).expect("test should finish"); + server.await.expect("server task should finish"); + } + + #[tokio::test] + async fn initial_connection_is_shared_by_all_waiters() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let server = tokio::spawn(async move { + let mut connection = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut connection, + "startup-session", + /*expected_resume_session_id*/ None, + ) + .await; + timeout(Duration::from_secs(1), connection.next()) + .await + .expect("client should close after the test"); + }); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + http_headers: HeaderMap::new(), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + + assert!(!client.startup_finished()); + let _startup_task = client.start_connecting(); + let (ready, first, second) = + tokio::join!(client.wait_until_ready(), client.get(), client.get()); + ready.expect("background startup should finish"); + let first = first.expect("first waiter should receive the client"); + let second = second.expect("second waiter should receive the same client"); + + assert!(client.startup_finished()); + assert_eq!(first.session_id().as_deref(), Some("startup-session")); + assert!(Arc::ptr_eq(&first.inner, &second.inner)); + + drop(first); + drop(second); + drop(client); + server.await.expect("server task should finish"); + } + + #[tokio::test] + async fn terminal_stdio_startup_failure_is_remembered() { + let client = LazyRemoteExecServerClient::new( + 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), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + + assert!(client.start_connecting().is_none()); + assert!(!client.startup_finished()); + let first = match client.get().await { + Ok(_) => panic!("missing executable should fail"), + Err(error) => error, + }; + assert!(client.startup_finished()); + let second = match client.get().await { + Ok(_) => panic!("burned environment should stay failed"), + Err(error) => error, + }; + assert!(matches!( + client.status().await, + EnvironmentObservedStatus::Disconnected { .. } + )); + + let ( + super::ExecServerError::ConnectionAttempt(first), + super::ExecServerError::ConnectionAttempt(second), + ) = (first, second) + else { + panic!("expected saved connection failures"); + }; + assert!(Arc::ptr_eq(&first, &second)); + } + + #[tokio::test] + async fn retryable_startup_failure_does_not_burn_environment() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let (replacement_initialized_tx, replacement_initialized_rx) = oneshot::channel(); + let server = tokio::spawn(async move { + let (mut failed_startup, _) = listener.accept().await.expect("startup should arrive"); + failed_startup + .write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n") + .await + .expect("failed handshake response should write"); + + let mut replacement = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut replacement, + "replacement-session", + /*expected_resume_session_id*/ None, + ) + .await; + replacement_initialized_tx + .send(()) + .expect("replacement initialization should be observed"); + timeout(Duration::from_secs(1), replacement.next()) + .await + .expect("client should close after the test"); + }); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + http_headers: HeaderMap::new(), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + + let failed_startup = match client.get().await { + Ok(_) => panic!("initial connection should fail"), + Err(error) => error, + }; + assert!(matches!( + failed_startup, + super::ExecServerError::ConnectionAttempt(_) + )); + assert!(client.startup_finished()); + assert!(matches!( + client.status().await, + EnvironmentObservedStatus::Disconnected { .. } + )); + + let (ready, first, second) = + tokio::join!(client.wait_until_ready(), client.get(), client.get()); + ready.expect("later readiness check should retry startup"); + let first = first.expect("first waiter should receive the replacement client"); + let second = second.expect("second waiter should receive the same replacement client"); + assert_eq!(first.session_id().as_deref(), Some("replacement-session")); + assert!(Arc::ptr_eq(&first.inner, &second.inner)); + replacement_initialized_rx + .await + .expect("server should observe replacement initialization"); + + drop(first); + drop(second); + drop(client); + server.await.expect("server task should finish"); + } + + #[tokio::test] + async fn failed_reconnect_does_not_burn_environment() { + let listener = TcpListener::bind("127.0.0.1:0") + .await + .expect("listener should bind"); + let websocket_url = format!( + "ws://{}", + listener.local_addr().expect("listener should have address") + ); + let (replacement_initialized_tx, replacement_initialized_rx) = oneshot::channel(); + let (allow_replacement_tx, allow_replacement_rx) = watch::channel(false); + let server = tokio::spawn(async move { + let mut first = accept_websocket(&listener).await; + complete_websocket_initialize( + &mut first, + "startup-session", + /*expected_resume_session_id*/ None, + ) + .await; + first + .close(None) + .await + .expect("startup websocket should close"); + + let successful_reconnect = loop { + let (stream, _) = listener.accept().await.expect("reconnect should arrive"); + if *allow_replacement_rx.borrow() { + break stream; + } + let mut failed_reconnect = stream; + failed_reconnect + .write_all(b"HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\n\r\n") + .await + .expect("failed handshake response should write"); + }; + let mut successful_reconnect = accept_async(successful_reconnect) + .await + .expect("replacement websocket handshake should succeed"); + complete_websocket_initialize( + &mut successful_reconnect, + "replacement-session", + /*expected_resume_session_id*/ None, + ) + .await; + replacement_initialized_tx + .send(()) + .expect("replacement initialization should be observed"); + timeout(Duration::from_secs(1), successful_reconnect.next()) + .await + .expect("client should close after the test"); + }); + let client = LazyRemoteExecServerClient::new( + ExecServerTransportParams::WebSocketUrl { + websocket_url, + connect_timeout: Duration::from_secs(1), + initialize_timeout: Duration::from_secs(1), + http_headers: HeaderMap::new(), + }, + HttpClientFactory::new(OutboundProxyPolicy::ReqwestDefault), + ); + + let initial = client.get().await.expect("startup should connect"); + timeout(Duration::from_secs(1), async { + while !initial.is_disconnected() { + tokio::task::yield_now().await; + } + }) + .await + .expect("client should observe disconnect"); + let failed_reconnect = match client.get().await { + Ok(_) => panic!("first lazy reconnect should fail"), + Err(error) => error, + }; + assert!(matches!( + failed_reconnect, + super::ExecServerError::ConnectionAttempt(_) + )); + allow_replacement_tx + .send(true) + .expect("server should allow a fresh client"); + let replacement = client.get().await.expect("later reconnect should succeed"); + + assert_eq!( + replacement.session_id().as_deref(), + Some("replacement-session") + ); + replacement_initialized_rx + .await + .expect("server should observe replacement initialization"); + + drop(initial); + drop(replacement); + drop(client); + server.await.expect("server task should finish"); + } + + #[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 { + session_id: "session-1".to_string(), + environment_info: None, + }) + .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, + sandbox_denied: Some(false), + }) + .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"); + } + + mod network_policy_tests; +}