codex-rs/exec-server/src/client_recovery.rs

use std::collections::hash_map::DefaultHasher;
use std::hash::Hash;
use std::hash::Hasher;
use std::sync::Arc;
use std::sync::atomic::Ordering;
use std::time::Duration;

use codex_network_proxy::NetworkDecision;
use codex_network_proxy::NetworkPolicyDecision;
use codex_network_proxy::NetworkPolicyRequest;
use codex_network_proxy::NetworkProtocol;
use codex_network_proxy::NetworkRequestCancellation;
use codex_network_proxy::NetworkRequestCancellationReason;
use serde_json::Value;
use tokio::sync::mpsc;
use tokio::time::Instant;
use tokio::time::sleep;
use tokio::time::timeout;
use tokio::time::timeout_at;
use tokio_util::sync::CancellationToken;
use tracing::Instrument;
use tracing::debug;

use super::ConnectionStatus;
use super::ExecServerClient;
use super::ExecServerError;
use super::Inner;
use super::OrderedSessionEvents;
use super::RecoveryPolicy;
use super::SessionState;
use super::disconnected_message;
use super::fail_all_in_flight_work;
use super::handle_server_notification;
use super::is_transport_closed_error;
use crate::client_transport::ExecServerReconnectStrategy;
use crate::process::ExecProcessEvent;
use crate::protocol::EXEC_READ_METHOD;
use crate::protocol::EXEC_TERMINATE_METHOD;
use crate::protocol::ExecServerNetworkPolicyDecision;
use crate::protocol::ExecServerNetworkProtocol;
use crate::protocol::MAX_NETWORK_POLICY_HOST_BYTES;
use crate::protocol::MAX_NETWORK_POLICY_PROCESS_ID_BYTES;
use crate::protocol::MAX_NETWORK_POLICY_REASON_BYTES;
use crate::protocol::NETWORK_POLICY_REQUEST_METHOD;
use crate::protocol::NetworkPolicyRequestParams;
use crate::protocol::NetworkPolicyRequestResponse;
use crate::protocol::ReadParams;
use crate::protocol::ReadResponse;
use crate::protocol::TerminateParams;
use crate::protocol::TerminateResponse;
use crate::rpc::RpcClient;
use crate::rpc::RpcClientEvent;
use crate::rpc::RpcInboundRequestAdmissionError;
use crate::rpc::SESSION_ALREADY_ATTACHED_ERROR_CODE;
use crate::rpc::invalid_params;
use crate::rpc::method_not_found;

#[cfg(test)]
const SESSION_RECOVERY_TIMEOUT: Duration = Duration::from_millis(500);
#[cfg(not(test))]
// Leave margin inside the server's 30-second retention windows because the
// client and server start their disconnect clocks independently.
const SESSION_RECOVERY_TIMEOUT: Duration = Duration::from_secs(25);
const SESSION_RECOVERY_RETRY_INTERVAL: Duration = Duration::from_millis(100);
const REGISTRY_RECOVERY_INITIAL_RETRY_INTERVAL: Duration = Duration::from_millis(500);
const REGISTRY_RECOVERY_MAX_RETRY_INTERVAL: Duration = Duration::from_secs(5);
const NETWORK_POLICY_DENIAL_REASON: &str = "not_allowed";

struct ClientRequestOutcome {
    span: tracing::Span,
    result: &'static str,
}

impl ClientRequestOutcome {
    fn complete(&mut self, result: &'static str) {
        self.result = result;
    }
}

impl Drop for ClientRequestOutcome {
    fn drop(&mut self) {
        self.span.record("result", self.result);
    }
}

impl SessionState {
    fn last_published_seq(&self) -> u64 {
        self.ordered_events
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner)
            .last_published_seq
    }

    fn recover_events(&self, response: ReadResponse) -> Result<bool, ExecServerError> {
        let ReadResponse {
            chunks,
            next_seq,
            exited,
            exit_code,
            closed,
            failure,
            sandbox_denied,
        } = response;
        if let Some(message) = failure {
            return Err(ExecServerError::Protocol(format!(
                "process failed while recovering: {message}"
            )));
        }

        let target_seq = next_seq.saturating_sub(1);
        let published_closed = {
            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
                || target_seq <= ordered_events.last_published_seq
            {
                return Ok(false);
            }
            let pending_exit = ordered_events.pending.range_mut(..=target_seq).find_map(
                |(_, event)| match event {
                    ExecProcessEvent::Exited {
                        sandbox_denied: pending_sandbox_denied,
                        ..
                    } => Some(pending_sandbox_denied),
                    _ => None,
                },
            );
            let exit_pending = pending_exit.is_some();
            if let Some(pending_sandbox_denied) = pending_exit {
                *pending_sandbox_denied =
                    Some(pending_sandbox_denied.unwrap_or(false) || sandbox_denied);
            }
            let mut exit_known = ordered_events.exit_published || exit_pending;
            if closed
                && (matches!(
                    ordered_events.pending.get(&target_seq),
                    Some(event) if !matches!(event, ExecProcessEvent::Closed { .. })
                ) || chunks.iter().any(|chunk| chunk.seq == target_seq))
            {
                return Err(ExecServerError::Protocol(format!(
                    "process close sequence {target_seq} conflicts with recovered output"
                )));
            }
            let mut published_closed = false;
            for chunk in chunks {
                if chunk.seq > target_seq {
                    return Err(ExecServerError::Protocol(format!(
                        "recovered process output sequence {} exceeds target sequence {target_seq}",
                        chunk.seq
                    )));
                }
                let next_seq = ordered_events.last_published_seq.saturating_add(1);
                if exited && !exit_known && chunk.seq > next_seq {
                    let exit_code = exit_code.ok_or_else(|| {
                        ExecServerError::Protocol(
                            "recovering exited process did not include its exit code".to_string(),
                        )
                    })?;
                    ordered_events
                        .insert_pending(ExecProcessEvent::Exited {
                            seq: next_seq,
                            exit_code,
                            sandbox_denied: Some(sandbox_denied),
                        })
                        .map_err(ExecServerError::Protocol)?;
                    published_closed |= self.publish_ready(&mut ordered_events);
                    exit_known = true;
                }
                if chunk.seq > ordered_events.last_published_seq {
                    ordered_events
                        .insert_pending(ExecProcessEvent::Output(chunk))
                        .map_err(ExecServerError::Protocol)?;
                    published_closed |= self.publish_ready(&mut ordered_events);
                }
            }
            if closed
                && !ordered_events.closed_published
                && !matches!(
                    ordered_events.pending.get(&target_seq),
                    Some(ExecProcessEvent::Closed { .. })
                )
            {
                ordered_events
                    .insert_pending(ExecProcessEvent::Closed { seq: target_seq })
                    .map_err(ExecServerError::Protocol)?;
            }

            let event_count = target_seq.saturating_sub(ordered_events.last_published_seq);
            let first_unpublished_seq = ordered_events.last_published_seq.saturating_add(1);
            let retained_count = if first_unpublished_seq <= target_seq {
                ordered_events
                    .pending
                    .range(first_unpublished_seq..=target_seq)
                    .count() as u64
            } else {
                0
            };
            let missing_count = event_count.saturating_sub(retained_count);
            if exited && !exit_known {
                if missing_count != 1 {
                    return Err(recovery_gap_error(target_seq));
                }
                let seq = first_missing_seq(&ordered_events, target_seq);
                let exit_code = exit_code.ok_or_else(|| {
                    ExecServerError::Protocol(
                        "recovering exited process did not include its exit code".to_string(),
                    )
                })?;
                ordered_events
                    .insert_pending(ExecProcessEvent::Exited {
                        seq,
                        exit_code,
                        sandbox_denied: Some(sandbox_denied),
                    })
                    .map_err(ExecServerError::Protocol)?;
            } else if missing_count != 0 {
                return Err(recovery_gap_error(target_seq));
            }
            published_closed |= self.publish_ready(&mut ordered_events);
            published_closed
        };

        self.note_change(target_seq);
        Ok(published_closed)
    }
}

fn first_missing_seq(events: &OrderedSessionEvents, target_seq: u64) -> u64 {
    let mut expected = events.last_published_seq.saturating_add(1);
    for seq in events
        .pending
        .range(expected..=target_seq)
        .map(|(seq, _)| *seq)
    {
        if seq != expected {
            break;
        }
        expected = expected.saturating_add(1);
    }
    expected
}

fn recovery_gap_error(target_seq: u64) -> ExecServerError {
    ExecServerError::Protocol(format!(
        "process events are no longer retained while recovering through sequence {target_seq}"
    ))
}

impl Inner {
    pub(super) async fn rpc_client(self: &Arc<Self>) -> Result<Arc<RpcClient>, ExecServerError> {
        let mut connection_changed = self.connection_changed.subscribe();
        loop {
            if let Some(message) = self.failure_message() {
                return Err(ExecServerError::Disconnected(message));
            }

            let rpc_client = {
                let connection = self
                    .connection
                    .lock()
                    .unwrap_or_else(std::sync::PoisonError::into_inner);
                match &connection.status {
                    ConnectionStatus::Connected(rpc_client) => Some(Arc::clone(rpc_client)),
                    ConnectionStatus::Recovering | ConnectionStatus::Failed(_) => None,
                }
            };
            let Some(rpc_client) = rpc_client else {
                let _ = connection_changed.changed().await;
                continue;
            };
            if !rpc_client.is_disconnected() {
                return Ok(rpc_client);
            }

            let _ = connection_changed.changed().await;
        }
    }

    pub(super) fn begin_process_start(&self, expected: &Arc<RpcClient>) -> bool {
        let mut connection = self
            .connection
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        let ConnectionStatus::Connected(current) = &connection.status else {
            return false;
        };
        if !Arc::ptr_eq(current, expected) || expected.is_disconnected() {
            return false;
        }
        connection.active_process_starts += 1;
        true
    }

    pub(super) fn finish_process_start(&self) {
        {
            let mut connection = self
                .connection
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            if connection.active_process_starts == 0 {
                tracing::error!("finished an exec-server process start that was not active");
                return;
            }
            connection.active_process_starts -= 1;
        }
        self.notify_connection_changed();
    }

    pub(super) fn is_failed(&self) -> bool {
        self.failure_message().is_some()
    }

    pub(super) fn failure_message(&self) -> Option<String> {
        let connection = self
            .connection
            .lock()
            .unwrap_or_else(std::sync::PoisonError::into_inner);
        match &connection.status {
            ConnectionStatus::Failed(message) => Some(message.clone()),
            ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => None,
        }
    }

    pub(super) fn request_recovery(
        self: &Arc<Self>,
        failed_rpc_client: Arc<RpcClient>,
        disconnect_message: String,
    ) {
        let should_recover = {
            let mut connection = self
                .connection
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            match &connection.status {
                ConnectionStatus::Connected(current)
                    if Arc::ptr_eq(current, &failed_rpc_client) =>
                {
                    connection.set_status(ConnectionStatus::Recovering);
                    true
                }
                ConnectionStatus::Connected(_)
                | ConnectionStatus::Recovering
                | ConnectionStatus::Failed(_) => false,
            }
        };
        if !should_recover {
            return;
        }

        self.notify_connection_changed();
        let inner = Arc::clone(self);
        tokio::spawn(async move {
            tokio::select! {
                biased;
                _ = inner.retired.cancelled() => {},
                _ = inner.recover(disconnect_message) => {},
            }
        });
    }

    async fn recover(self: &Arc<Self>, disconnect_message: String) {
        let deadline = Instant::now() + SESSION_RECOVERY_TIMEOUT;
        self.fail_all_http_body_streams(disconnect_message.clone())
            .await;
        if timeout_at(deadline, self.wait_for_process_starts())
            .await
            .is_err()
        {
            let message = format!(
                "{disconnect_message}; failed to resume exec-server session: recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}"
            );
            self.fail(message).await;
            return;
        }
        if self.reconnect_strategy.is_none() {
            self.fail(disconnect_message).await;
            return;
        }

        let Some(session_id) = self.session_id.get().cloned() else {
            let message = format!(
                "{disconnect_message}; failed to resume exec-server session: missing session id"
            );
            self.fail(message).await;
            return;
        };
        let uses_registry_backoff = matches!(
            self.reconnect_strategy.as_ref(),
            Some(ExecServerReconnectStrategy::NoiseRendezvous { .. })
        );
        let mut registry_retry_attempt = 0;
        let last_error = loop {
            match timeout_at(deadline, self.resume_once(&session_id)).await {
                Ok(Ok((rpc_client, _attempt))) => {
                    if !rpc_client.is_disconnected() && self.install_recovered_client(rpc_client) {
                        return;
                    }
                }
                Ok(Err(error)) if !is_retryable_recovery_error(&error) => {
                    break error.to_string();
                }
                Ok(Err(_)) => {}
                Err(_) => {
                    break format!("recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}");
                }
            }

            let retry_delay = if uses_registry_backoff {
                let delay = registry_recovery_retry_delay(&session_id, registry_retry_attempt);
                registry_retry_attempt = registry_retry_attempt.saturating_add(1);
                delay
            } else {
                SESSION_RECOVERY_RETRY_INTERVAL
            };

            let now = Instant::now();
            if now >= deadline {
                break format!("recovery timed out after {SESSION_RECOVERY_TIMEOUT:?}");
            }
            sleep(retry_delay.min(deadline - now)).await;
        };

        let message =
            format!("{disconnect_message}; failed to resume exec-server session: {last_error}");
        self.fail(message).await;
    }

    async fn wait_for_process_starts(&self) {
        let mut connection_changed = self.connection_changed.subscribe();
        loop {
            let starts_are_done = self
                .connection
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner)
                .active_process_starts
                == 0;
            if starts_are_done {
                return;
            }
            let _ = connection_changed.changed().await;
        }
    }

    fn install_recovered_client(&self, rpc_client: Arc<RpcClient>) -> bool {
        let installed = {
            let mut connection = self
                .connection
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            if !matches!(connection.status, ConnectionStatus::Recovering)
                || rpc_client.is_disconnected()
            {
                false
            } else {
                connection.set_status(ConnectionStatus::Connected(rpc_client));
                true
            }
        };
        if installed {
            self.notify_connection_changed();
        }
        installed
    }

    fn notify_connection_changed(&self) {
        self.connection_changed.send_replace(());
    }

    async fn resume_once(
        self: &Arc<Self>,
        session_id: &str,
    ) -> Result<(Arc<RpcClient>, Option<tokio::sync::OwnedSemaphorePermit>), ExecServerError> {
        let reconnect_strategy = self
            .reconnect_strategy
            .as_ref()
            .ok_or_else(|| ExecServerError::Protocol("missing reconnect strategy".to_string()))?;
        let attempt = reconnect_strategy.resume(session_id).await?;
        let (connection, options, attempt_permit, noise_context) = attempt.into_parts();
        let (rpc_client, events_rx) = RpcClient::new(connection);
        let rpc_client = Arc::new(rpc_client);
        let client = ExecServerClient {
            inner: Arc::clone(self),
            recovery_policy: RecoveryPolicy::Wait,
        };
        // Resuming a session redirects notifications from its running processes
        // to this connection during initialize. Drain them immediately so a
        // burst cannot fill the bounded event channel and block the initialize
        // response behind it.
        client.spawn_rpc_reader(&rpc_client, events_rx);
        client
            .initialize_rpc(&rpc_client, options, noise_context)
            .await?;

        self.recover_processes(&rpc_client).await?;
        Ok((rpc_client, attempt_permit))
    }

    async fn recover_processes(
        self: &Arc<Self>,
        rpc_client: &RpcClient,
    ) -> Result<(), ExecServerError> {
        let sessions = self.sessions.load_full();
        for (process_id, session) in sessions.iter() {
            if !session.recoverable.load(Ordering::Acquire) {
                continue;
            }
            let response = rpc_client
                .call::<_, ReadResponse>(
                    EXEC_READ_METHOD,
                    &ReadParams {
                        process_id: process_id.clone(),
                        after_seq: Some(session.last_published_seq()),
                        max_bytes: None,
                        wait_ms: Some(0),
                    },
                )
                .await
                .map_err(ExecServerError::from);
            let recovered = match response {
                Ok(response) => session.recover_events(response),
                Err(error) if is_transport_closed_error(&error) => return Err(error),
                Err(error) => Err(error),
            };
            match recovered {
                Ok(true) => self.remove_session_if(process_id, session),
                Ok(false) => {}
                Err(error) => {
                    session
                        .network_policy
                        .cancellation
                        .record(NetworkRequestCancellationReason::ProcessCancelled);
                    let terminated: Result<TerminateResponse, ExecServerError> = rpc_client
                        .call_for_cleanup(
                            EXEC_TERMINATE_METHOD,
                            &TerminateParams {
                                process_id: process_id.clone(),
                            },
                        )
                        .await
                        .map_err(ExecServerError::from);
                    if let Err(terminate_error) = terminated
                        && is_transport_closed_error(&terminate_error)
                    {
                        return Err(terminate_error);
                    }
                    self.remove_session_if(process_id, session);
                    session.set_failure(format!("failed to recover process {process_id}: {error}"));
                }
            }
        }
        Ok(())
    }

    async fn fail(self: &Arc<Self>, message: String) {
        let (message, newly_failed) = {
            let mut connection = self
                .connection
                .lock()
                .unwrap_or_else(std::sync::PoisonError::into_inner);
            match &connection.status {
                ConnectionStatus::Failed(existing) => (existing.clone(), false),
                ConnectionStatus::Connected(_) | ConnectionStatus::Recovering => {
                    connection.set_status(ConnectionStatus::Failed(message.clone()));
                    (message, true)
                }
            }
        };
        if newly_failed {
            self.notify_connection_changed();
            fail_all_in_flight_work(self, message.clone()).await;
        }
    }
}

impl ExecServerClient {
    pub(super) fn spawn_rpc_reader(
        &self,
        rpc_client: &Arc<RpcClient>,
        mut events_rx: mpsc::Receiver<RpcClientEvent>,
    ) {
        let inner = Arc::downgrade(&self.inner);
        let rpc_inbound_request_slots = Arc::clone(&self.inner.rpc_inbound_request_slots);
        let rpc_client = Arc::downgrade(rpc_client);
        let connection_cancelled = CancellationToken::new();
        let connection_cancel_guard = connection_cancelled.clone().drop_guard();
        tokio::spawn(async move {
            let _connection_cancel_guard = connection_cancel_guard;
            while let Some(event) = events_rx.recv().await {
                let (Some(inner), Some(rpc_client)) = (inner.upgrade(), rpc_client.upgrade())
                else {
                    return;
                };
                match event {
                    RpcClientEvent::Request {
                        request,
                        request_span,
                    } => {
                        let mut request_outcome = ClientRequestOutcome {
                            span: request_span,
                            result: "disconnected",
                        };
                        if request.method != NETWORK_POLICY_REQUEST_METHOD {
                            let error = method_not_found(format!(
                                "exec-server client does not implement `{}` yet",
                                request.method
                            ));
                            if rpc_client.respond_error(request.id, error).await.is_err() {
                                inner.request_recovery(
                                    rpc_client,
                                    disconnected_message(/*reason*/ None),
                                );
                                return;
                            }
                            request_outcome.complete("error");
                            continue;
                        }
                        request_outcome
                            .span
                            .record("otel.name", NETWORK_POLICY_REQUEST_METHOD);

                        let request_guard = match rpc_client
                            .admit_inbound_request(&request.id, &rpc_inbound_request_slots)
                        {
                            Ok(request_guard) => request_guard,
                            Err(RpcInboundRequestAdmissionError::InvalidRequestId) => {
                                rpc_client.close_transport().await;
                                inner.request_recovery(
                                    rpc_client,
                                    "exec-server sent an invalid request ID".to_string(),
                                );
                                return;
                            }
                            Err(RpcInboundRequestAdmissionError::DuplicateRequestId) => {
                                rpc_client.close_transport().await;
                                inner.request_recovery(
                                    rpc_client,
                                    "exec-server reused an in-flight request ID".to_string(),
                                );
                                return;
                            }
                            Err(RpcInboundRequestAdmissionError::AtCapacity) => {
                                let response = NetworkPolicyRequestResponse {
                                    decision: ExecServerNetworkPolicyDecision::Deny {
                                        reason: NETWORK_POLICY_DENIAL_REASON.to_string(),
                                    },
                                };
                                if rpc_client.respond(request.id, &response).await.is_err() {
                                    inner.request_recovery(
                                        rpc_client,
                                        disconnected_message(/*reason*/ None),
                                    );
                                    return;
                                }
                                request_outcome.complete("success");
                                continue;
                            }
                        };
                        let request_id = request.id;
                        let params: NetworkPolicyRequestParams =
                            match serde_json::from_value(request.params.unwrap_or(Value::Null)) {
                                Ok(params) => params,
                                Err(_) => {
                                    let error = invalid_params(
                                        "invalid network policy request params".to_string(),
                                    );
                                    if rpc_client.respond_error(request_id, error).await.is_err() {
                                        inner.request_recovery(
                                            rpc_client,
                                            disconnected_message(/*reason*/ None),
                                        );
                                        return;
                                    }
                                    request_outcome.complete("error");
                                    continue;
                                }
                            };
                        let process_id = params.process_id;
                        let request = params.request;
                        let process_id_valid = !process_id.is_empty()
                            && process_id.len() <= MAX_NETWORK_POLICY_PROCESS_ID_BYTES;
                        let host_valid = !request.host.is_empty()
                            && request.host.len() <= MAX_NETWORK_POLICY_HOST_BYTES
                            && !request.host.chars().any(char::is_control)
                            && !request.host.chars().any(char::is_whitespace);
                        let session = (process_id_valid && host_valid)
                            .then(|| inner.get_session(&process_id))
                            .flatten();
                        let controller = session
                            .as_ref()
                            .and_then(|session| session.network_policy.controller.load_full());
                        let process_cancelled = session
                            .as_ref()
                            .map(|session| session.network_policy.cancelled.clone());
                        let process_cancellation = session
                            .as_ref()
                            .map(|session| session.network_policy.cancellation.clone());
                        let cancellation = NetworkRequestCancellation::default();
                        let expected_session = session.as_ref().map(Arc::downgrade);
                        let policy_request =
                            (process_id_valid && host_valid).then_some(NetworkPolicyRequest {
                                protocol: match request.protocol {
                                    ExecServerNetworkProtocol::Http => NetworkProtocol::Http,
                                    ExecServerNetworkProtocol::HttpsConnect => {
                                        NetworkProtocol::HttpsConnect
                                    }
                                    ExecServerNetworkProtocol::Socks5Tcp => {
                                        NetworkProtocol::Socks5Tcp
                                    }
                                    ExecServerNetworkProtocol::Socks5Udp => {
                                        NetworkProtocol::Socks5Udp
                                    }
                                },
                                host: request.host,
                                port: request.port,
                                environment_id: None,
                                client_addr: None,
                                method: None,
                                command: None,
                                exec_policy_hint: None,
                                execution_id: None,
                                disconnect: None,
                                cancellation: Some(cancellation.clone()),
                            });
                        let inner = Arc::downgrade(&inner);
                        let rpc_client = Arc::downgrade(&rpc_client);
                        let connection_cancelled = connection_cancelled.clone();
                        let task_span = request_outcome.span.clone();
                        let task = async move {
                            let _request_guard = request_guard;
                            let decision = match (controller, policy_request, process_cancelled) {
                                (Some(controller), Some(request), Some(process_cancelled)) => {
                                    // Keep the decision future outside select/timeout so its
                                    // guard sees the cancellation cause before it is dropped.
                                    let mut decision = controller.decider.decide(request);
                                    tokio::select! {
                                        biased;
                                        _ = connection_cancelled.cancelled() => {
                                            cancellation.record(NetworkRequestCancellationReason::ConnectionClosed);
                                            return;
                                        },
                                        _ = process_cancelled.cancelled() => {
                                            cancellation.record(process_cancellation.as_ref()
                                                .and_then(NetworkRequestCancellation::reason)
                                                .unwrap_or(NetworkRequestCancellationReason::ProcessCancelled));
                                            NetworkDecision::deny(NETWORK_POLICY_DENIAL_REASON)
                                        }
                                        result = timeout(
                                            controller.timeout,
                                            &mut decision,
                                        ) => result.unwrap_or_else(|_| {
                                            cancellation.record(NetworkRequestCancellationReason::TimedOut);
                                            NetworkDecision::deny(NETWORK_POLICY_DENIAL_REASON)
                                        }),
                                    }
                                }
                                (None, _, _) | (_, None, _) | (_, _, None) => {
                                    NetworkDecision::deny(NETWORK_POLICY_DENIAL_REASON)
                                }
                            };
                            if let Some(expected_session) = expected_session {
                                let (Some(inner), Some(expected_session)) =
                                    (inner.upgrade(), expected_session.upgrade())
                                else {
                                    return;
                                };
                                if !inner
                                    .get_session(&process_id)
                                    .is_some_and(|session| Arc::ptr_eq(&session, &expected_session))
                                {
                                    return;
                                }
                            }
                            let Some(rpc_client) = rpc_client.upgrade() else {
                                return;
                            };
                            let decision = match decision {
                                NetworkDecision::Allow => ExecServerNetworkPolicyDecision::Allow,
                                NetworkDecision::Deny {
                                    reason, decision, ..
                                } if reason.len() <= MAX_NETWORK_POLICY_REASON_BYTES
                                    && !reason.chars().any(char::is_control) =>
                                {
                                    match decision {
                                        NetworkPolicyDecision::Deny => {
                                            ExecServerNetworkPolicyDecision::Deny { reason }
                                        }
                                        NetworkPolicyDecision::Ask => {
                                            ExecServerNetworkPolicyDecision::Ask { reason }
                                        }
                                    }
                                }
                                NetworkDecision::Deny { .. } => {
                                    ExecServerNetworkPolicyDecision::Deny {
                                        reason: NETWORK_POLICY_DENIAL_REASON.to_string(),
                                    }
                                }
                            };
                            if let Err(error) = rpc_client
                                .respond(request_id, &NetworkPolicyRequestResponse { decision })
                                .await
                            {
                                debug!(
                                    ?error,
                                    "failed to send network policy decision to exec-server"
                                );
                            } else {
                                request_outcome.complete("success");
                            }
                        };
                        tokio::spawn(task.instrument(task_span));
                    }
                    RpcClientEvent::Notification(notification) => {
                        if let Err(error) = handle_server_notification(&inner, notification).await {
                            rpc_client.close_transport().await;
                            inner.request_recovery(
                                rpc_client,
                                format!("exec-server notification handling failed: {error}"),
                            );
                            return;
                        }
                    }
                    RpcClientEvent::Disconnected { reason } => {
                        inner.request_recovery(rpc_client, disconnected_message(reason.as_deref()));
                        return;
                    }
                }
            }
        });
    }
}

pub(crate) fn is_retryable_recovery_error(error: &ExecServerError) -> bool {
    if let ExecServerError::ConnectionAttempt(error) = error {
        return is_retryable_recovery_error(error.as_ref());
    }
    is_transport_closed_error(error)
        || matches!(
            error,
            ExecServerError::ProvisioningFailed(_)
                | ExecServerError::WebSocketConnectTimeout { .. }
                | ExecServerError::WebSocketConnect { .. }
                | ExecServerError::InitializeTimedOut { .. }
        )
        || is_retryable_registry_error(error)
        || matches!(
            error,
            ExecServerError::Server { code, .. }
                if *code == SESSION_ALREADY_ATTACHED_ERROR_CODE
        )
}

pub(crate) fn is_retryable_registry_error(error: &ExecServerError) -> bool {
    matches!(
        error,
        ExecServerError::EnvironmentRegistryRequest(error)
            if error.is_connect()
                || error.is_timeout()
                || error.is_body()
                || matches!(
                    error,
                    codex_http_client::RouteAwareRequestError::Request(error)
                        if error.is_decode()
                )
    ) || matches!(
        error,
        ExecServerError::EnvironmentRegistryHttp { status, .. }
            if status.is_server_error()
                || *status == http::StatusCode::REQUEST_TIMEOUT
                || *status == http::StatusCode::TOO_MANY_REQUESTS
    ) || is_environment_offline_error(error)
}

pub(crate) fn is_environment_offline_error(error: &ExecServerError) -> bool {
    matches!(
        error,
        ExecServerError::EnvironmentRegistryHttp { status, code, .. }
            if *status == http::StatusCode::CONFLICT
                && code.as_deref() == Some("environment_offline")
    )
}

pub(crate) fn registry_recovery_retry_delay(retry_key: &str, attempt: u32) -> Duration {
    let multiplier = 1_u32.checked_shl(attempt.min(4)).unwrap_or(u32::MAX);
    let base_delay = REGISTRY_RECOVERY_INITIAL_RETRY_INTERVAL
        .saturating_mul(multiplier)
        .min(REGISTRY_RECOVERY_MAX_RETRY_INTERVAL);
    let base_millis = base_delay.as_millis() as u64;
    let mut hasher = DefaultHasher::new();
    retry_key.hash(&mut hasher);
    attempt.hash(&mut hasher);

    Duration::from_millis(base_millis + hasher.finish() % (base_millis / 2 + 1))
}

#[cfg(test)]
#[path = "client_recovery_tests.rs"]
mod tests;