use super::CurrentRemoteControlEnrollment;
use super::RemoteControlEnrollmentSelection;
use super::RemoteControlPairingPersistenceKey;
use super::desired_state::RemoteControlDesiredState;
use super::desired_state::acquire_persistence_lock;
use super::desired_state::desired_state_from_persisted_enrollment;
use super::protocol::ClientEnvelope;
use super::protocol::ClientEvent;
use super::protocol::ClientId;
use super::protocol::RemoteControlTarget;
use super::protocol::ServerEnvelope;
use super::protocol::StreamId;
use super::remote_control_status_with_connection_status;
use super::same_remote_control_enrollment;
use super::segment::ClientSegmentObservation;
use super::segment::ClientSegmentReassembler;
use super::segment::REMOTE_CONTROL_SEGMENT_MAX_BYTES;
use super::segment::split_server_envelope_for_transport;
use crate::transport::TransportEvent;
use crate::transport::remote_control::auth::RemoteControlConnectionAuth;
use crate::transport::remote_control::auth::load_remote_control_auth;
use crate::transport::remote_control::auth::recover_remote_control_auth;
use crate::transport::remote_control::client_tracker::ClientTracker;
use crate::transport::remote_control::client_tracker::REMOTE_CONTROL_IDLE_SWEEP_INTERVAL;
use crate::transport::remote_control::enroll::RemoteControlEnrollment;
use crate::transport::remote_control::enroll::enroll_remote_control_server;
use crate::transport::remote_control::enroll::format_headers;
use crate::transport::remote_control::enroll::load_persisted_remote_control_enrollment;
use crate::transport::remote_control::enroll::preview_remote_control_response_body;
use crate::transport::remote_control::enroll::refresh_remote_control_server;
use crate::transport::remote_control::enroll::update_persisted_remote_control_enrollment;
use axum::http::HeaderValue;
use base64::Engine;
use codex_app_server_protocol::RemoteControlConnectionStatus;
use codex_app_server_protocol::RemoteControlStatusChangedNotification;
use codex_core::util::backoff;
use codex_login::AuthManager;
use codex_login::UnauthorizedRecovery;
use codex_state::StateRuntime;
use codex_utils_rustls_provider::ensure_rustls_crypto_provider;
use futures::SinkExt;
use futures::StreamExt;
use futures::stream::SplitSink;
use futures::stream::SplitStream;
use std::collections::HashMap;
use std::collections::VecDeque;
use std::io;
use std::io::ErrorKind;
use std::sync::Arc;
use tokio::net::TcpStream;
use tokio::sync::Mutex;
use tokio::sync::Semaphore;
use tokio::sync::mpsc;
use tokio::sync::oneshot;
use tokio::sync::watch;
use tokio::time::MissedTickBehavior;
use tokio_tungstenite::MaybeTlsStream;
use tokio_tungstenite::WebSocketStream;
use tokio_tungstenite::connect_async;
use tokio_tungstenite::tungstenite;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_util::sync::CancellationToken;

#[cfg(test)]
use super::RemoteControlEnrollmentState;
use tracing::error;
use tracing::info;
use tracing::warn;

pub(super) const REMOTE_CONTROL_PROTOCOL_VERSION: &str = "3";
pub(super) const REMOTE_CONTROL_INSTALLATION_ID_HEADER: &str = "x-codex-installation-id";
const REMOTE_CONTROL_SUBSCRIBE_CURSOR_HEADER: &str = "x-codex-subscribe-cursor";
const REMOTE_CONTROL_WEBSOCKET_PING_INTERVAL: std::time::Duration =
    std::time::Duration::from_secs(10);
const REMOTE_CONTROL_WEBSOCKET_PONG_TIMEOUT: std::time::Duration =
    std::time::Duration::from_secs(60);
const REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL: std::time::Duration =
    std::time::Duration::from_secs(1);
const REMOTE_CONTROL_RECONNECT_BACKOFF_CAP: std::time::Duration =
    std::time::Duration::from_secs(30);
const REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT: std::time::Duration =
    std::time::Duration::from_secs(30);
const REMOTE_CONTROL_CONNECTION_SHUTDOWN_TIMEOUT: std::time::Duration =
    std::time::Duration::from_secs(5);
const REMOTE_APP_SERVER_NOT_FOUND_DETAIL: &str = "Remote app server not found";

struct BoundedOutboundBuffer {
    buffer_by_stream: HashMap<(ClientId, StreamId), VecDeque<ServerEnvelope>>,
    used_tx: watch::Sender<usize>,
}

impl BoundedOutboundBuffer {
    fn new() -> (Self, watch::Receiver<usize>) {
        let (used_tx, used_rx) = watch::channel(0);
        let buffer = Self {
            buffer_by_stream: HashMap::new(),
            used_tx,
        };
        (buffer, used_rx)
    }

    fn insert(&mut self, server_envelope: &ServerEnvelope) {
        self.buffer_by_stream
            .entry((
                server_envelope.client_id.clone(),
                server_envelope.stream_id.clone(),
            ))
            .or_default()
            .push_back(server_envelope.clone());
        self.used_tx.send_modify(|used| *used += 1);
    }

    fn ack(
        &mut self,
        client_id: &ClientId,
        stream_id: &StreamId,
        acked_seq_id: u64,
        acked_segment_id: Option<usize>,
    ) {
        let key = (client_id.clone(), stream_id.clone());
        let Some(buffer) = self.buffer_by_stream.get_mut(&key) else {
            return;
        };
        let acked_cursor = (acked_seq_id, acked_segment_id.unwrap_or(usize::MAX));
        buffer.retain(|server_envelope| {
            let envelope_cursor = (
                server_envelope.seq_id,
                server_envelope.event.segment_id().unwrap_or_default(),
            );
            let is_acked = envelope_cursor <= acked_cursor;
            if is_acked {
                self.used_tx.send_modify(|used| *used -= 1);
            }
            !is_acked
        });
        if buffer.is_empty() {
            self.buffer_by_stream.remove(&key);
        }
    }

    fn server_envelopes(&self) -> impl Iterator<Item = &ServerEnvelope> {
        self.buffer_by_stream
            .values()
            .flat_map(|buffer| buffer.iter())
    }
}

struct WebsocketState {
    outbound_buffer: BoundedOutboundBuffer,
    subscribe_cursor: Option<String>,
    next_seq_id_by_stream: HashMap<(ClientId, StreamId), u64>,
    last_completed_client_chunk_seq_id_by_stream: HashMap<(ClientId, Option<StreamId>), u64>,
    client_segment_reassembler: ClientSegmentReassembler,
}

impl WebsocketState {
    fn observe_client_message(
        &mut self,
        client_envelope: ClientEnvelope,
        wire_size_bytes: usize,
    ) -> ClientSegmentObservation {
        let client_message_key = Self::client_message_key(&client_envelope);
        if let Some((key, seq_id)) = client_message_key.as_ref()
            && self
                .last_completed_client_chunk_seq_id_by_stream
                .get(key)
                .is_some_and(|last_seq_id| last_seq_id >= seq_id)
        {
            return ClientSegmentObservation::Dropped;
        }
        if let (
            Some((_, seq_id)),
            Some(stream_id),
            ClientEvent::ClientMessageChunk { segment_id, .. },
        ) = (
            client_message_key.as_ref(),
            client_envelope.stream_id.as_ref(),
            &client_envelope.event,
        ) && self.client_segment_reassembler.should_ignore_chunk(
            &client_envelope.client_id,
            stream_id,
            *seq_id,
            *segment_id,
        ) {
            return ClientSegmentObservation::Dropped;
        }
        if client_message_key.is_some() && wire_size_bytes > REMOTE_CONTROL_SEGMENT_MAX_BYTES {
            warn!(
                client_id = client_envelope.client_id.0.as_str(),
                "dropping oversized segmented remote-control client envelope"
            );
            if let Some(stream_id) = client_envelope.stream_id.as_ref() {
                self.client_segment_reassembler
                    .invalidate_stream(&client_envelope.client_id, stream_id);
            }
            return ClientSegmentObservation::Dropped;
        }

        self.client_segment_reassembler.observe(client_envelope)
    }

    fn record_client_message_delivery(
        &mut self,
        client_envelope: &ClientEnvelope,
        client_message_key: Option<((ClientId, Option<StreamId>), u64)>,
    ) {
        if let Some(cursor) = client_envelope.cursor.as_deref() {
            self.subscribe_cursor = Some(cursor.to_string());
        }
        if let Some((key, seq_id)) = client_message_key {
            self.last_completed_client_chunk_seq_id_by_stream
                .insert(key, seq_id);
        }
        if let ClientEvent::Ack { segment_id } = &client_envelope.event
            && let Some(acked_seq_id) = client_envelope.seq_id
            && let Some(stream_id) = client_envelope.stream_id.as_ref()
        {
            self.outbound_buffer.ack(
                &client_envelope.client_id,
                stream_id,
                acked_seq_id,
                *segment_id,
            );
        }
    }

    fn invalidate_client_message_stream(&mut self, client_id: &ClientId, stream_id: &StreamId) {
        self.last_completed_client_chunk_seq_id_by_stream
            .remove(&(client_id.clone(), Some(stream_id.clone())));
    }

    fn invalidate_client_message_client(&mut self, client_id: &ClientId) {
        self.last_completed_client_chunk_seq_id_by_stream
            .retain(|(cursor_client_id, _), _| cursor_client_id != client_id);
    }

    fn client_message_key(
        client_envelope: &ClientEnvelope,
    ) -> Option<((ClientId, Option<StreamId>), u64)> {
        let seq_id = match (&client_envelope.event, client_envelope.seq_id) {
            (ClientEvent::ClientMessageChunk { .. }, Some(seq_id)) => seq_id,
            _ => return None,
        };
        Some((
            (
                client_envelope.client_id.clone(),
                client_envelope.stream_id.clone(),
            ),
            seq_id,
        ))
    }
}

pub(crate) struct RemoteControlWebsocket {
    remote_control_url: String,
    installation_id: String,
    server_name: String,
    remote_control_target: Option<RemoteControlTarget>,
    state_db: Option<Arc<StateRuntime>>,
    auth_manager: Arc<AuthManager>,
    status_publisher: RemoteControlStatusPublisher,
    shutdown_token: CancellationToken,
    reconnect_attempt: u64,
    auth_recovery: UnauthorizedRecovery,
    auth_change_rx: watch::Receiver<u64>,
    current_enrollment: CurrentRemoteControlEnrollment,
    pairing_persistence_key: RemoteControlPairingPersistenceKey,
    client_tracker: Arc<Mutex<ClientTracker>>,
    state: Arc<Mutex<WebsocketState>>,
    server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>,
    used_rx: watch::Receiver<usize>,
    desired_state_tx: Arc<watch::Sender<RemoteControlDesiredState>>,
    desired_state_rx: watch::Receiver<RemoteControlDesiredState>,
    desired_state_persistence_lock: Arc<Semaphore>,
}

pub(crate) struct RemoteControlWebsocketConfig {
    pub(crate) remote_control_url: String,
    pub(crate) installation_id: String,
    pub(crate) remote_control_target: Option<RemoteControlTarget>,
    pub(crate) server_name: String,
}

pub(super) struct RemoteControlAuthContext<'a> {
    auth_manager: &'a Arc<AuthManager>,
    auth_recovery: &'a mut UnauthorizedRecovery,
    auth_change_rx: &'a mut watch::Receiver<u64>,
}

struct RemoteControlEnrollmentAuthContext<'a, 'b> {
    auth: &'a RemoteControlConnectionAuth,
    recovery: &'a mut RemoteControlAuthContext<'b>,
}

enum ConnectOutcome {
    Connected(Box<WebSocketStream<MaybeTlsStream<TcpStream>>>),
    Disabled,
    Shutdown,
}

#[derive(Debug, Clone, Copy)]
enum ConnectionEndReason {
    Shutdown,
    Disabled,
    EnabledWatchClosed,
    ConnectionWorkerStopped,
}

pub(super) struct RemoteControlChannels {
    pub(super) transport_event_tx: mpsc::Sender<TransportEvent>,
    pub(super) status_publisher: RemoteControlStatusPublisher,
    pub(super) current_enrollment: CurrentRemoteControlEnrollment,
    pub(super) pairing_persistence_key: RemoteControlPairingPersistenceKey,
    pub(super) desired_state_persistence_lock: Arc<Semaphore>,
}

#[derive(Clone)]
pub(super) struct RemoteControlStatusPublisher {
    tx: watch::Sender<RemoteControlStatusChangedNotification>,
}

impl RemoteControlStatusPublisher {
    pub(super) fn new(tx: watch::Sender<RemoteControlStatusChangedNotification>) -> Self {
        Self { tx }
    }

    fn status(&self) -> RemoteControlStatusChangedNotification {
        self.tx.borrow().clone()
    }

    fn publish_status(&self, connection_status: RemoteControlConnectionStatus) {
        let mut status_change = None;
        self.tx.send_if_modified(|status| {
            let next_status =
                remote_control_status_with_connection_status(status, connection_status);
            if *status == next_status {
                return false;
            }

            status_change = Some((status.clone(), next_status.clone()));
            *status = next_status;
            true
        });
        if let Some((previous_status, next_status)) = status_change {
            info!(
                previous_status = ?previous_status.status,
                next_status = ?next_status.status,
                previous_environment_id = ?previous_status.environment_id,
                next_environment_id = ?next_status.environment_id,
                installation_id = %next_status.installation_id,
                server_name = %next_status.server_name,
                "remote control websocket status changed"
            );
        }
    }

    pub(super) fn publish_environment_id(&self, environment_id: Option<String>) {
        let mut status_change = None;
        self.tx.send_if_modified(|status| {
            if status.status == RemoteControlConnectionStatus::Disabled {
                return false;
            }
            let next_status = RemoteControlStatusChangedNotification {
                status: status.status,
                server_name: status.server_name.clone(),
                installation_id: status.installation_id.clone(),
                environment_id,
            };
            if *status == next_status {
                return false;
            }

            status_change = Some((status.clone(), next_status.clone()));
            *status = next_status;
            true
        });
        if let Some((previous_status, next_status)) = status_change {
            info!(
                status = ?next_status.status,
                previous_environment_id = ?previous_status.environment_id,
                next_environment_id = ?next_status.environment_id,
                installation_id = %next_status.installation_id,
                server_name = %next_status.server_name,
                "remote control websocket environment changed"
            );
        }
    }
}

#[derive(Clone, Copy)]
pub(super) struct RemoteControlConnectOptions<'a> {
    installation_id: &'a str,
    server_name: &'a str,
    subscribe_cursor: Option<&'a str>,
    app_server_client_name: Option<&'a str>,
    desired_state_tx: &'a watch::Sender<RemoteControlDesiredState>,
    desired_state_persistence_lock: &'a Semaphore,
}

impl RemoteControlWebsocket {
    pub(crate) fn new(
        config: RemoteControlWebsocketConfig,
        state_db: Option<Arc<StateRuntime>>,
        auth_manager: Arc<AuthManager>,
        channels: RemoteControlChannels,
        shutdown_token: CancellationToken,
        desired_state_tx: Arc<watch::Sender<RemoteControlDesiredState>>,
    ) -> Self {
        let shutdown_token = shutdown_token.child_token();
        let (server_event_tx, server_event_rx) = mpsc::channel(super::CHANNEL_CAPACITY);
        let client_tracker = ClientTracker::new(
            server_event_tx,
            channels.transport_event_tx,
            &shutdown_token,
        );
        let (outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
        let auth_recovery = auth_manager.unauthorized_recovery();
        let auth_change_rx = auth_manager.auth_change_receiver();

        let desired_state_rx = desired_state_tx.subscribe();
        Self {
            remote_control_url: config.remote_control_url,
            installation_id: config.installation_id,
            server_name: config.server_name,
            remote_control_target: config.remote_control_target,
            state_db,
            auth_manager,
            status_publisher: channels.status_publisher,
            shutdown_token,
            reconnect_attempt: 0,
            auth_recovery,
            auth_change_rx,
            current_enrollment: channels.current_enrollment,
            pairing_persistence_key: channels.pairing_persistence_key,
            client_tracker: Arc::new(Mutex::new(client_tracker)),
            state: Arc::new(Mutex::new(WebsocketState {
                outbound_buffer,
                subscribe_cursor: None,
                next_seq_id_by_stream: HashMap::new(),
                last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
                client_segment_reassembler: ClientSegmentReassembler::default(),
            })),
            server_event_rx: Arc::new(Mutex::new(server_event_rx)),
            used_rx,
            desired_state_tx,
            desired_state_rx,
            desired_state_persistence_lock: channels.desired_state_persistence_lock,
        }
    }

    #[expect(
        clippy::await_holding_invalid_type,
        reason = "remote-control client shutdown must serialize tracker state"
    )]
    pub(crate) async fn run(
        mut self,
        app_server_client_name_rx: Option<oneshot::Receiver<String>>,
    ) {
        info!(
            remote_control_url = %self.remote_control_url,
            installation_id = %self.installation_id,
            server_name = %self.server_name,
            initial_desired_state = ?*self.desired_state_rx.borrow(),
            "app-server remote control websocket loop started"
        );
        let app_server_client_name = match self
            .wait_for_app_server_client_name(app_server_client_name_rx)
            .await
        {
            Ok(app_server_client_name) => app_server_client_name,
            Err(_) => {
                warn!(
                    remote_control_url = %self.remote_control_url,
                    installation_id = %self.installation_id,
                    server_name = %self.server_name,
                    shutdown_requested = self.shutdown_token.is_cancelled(),
                    "app-server remote control websocket loop stopped before client name was ready"
                );
                self.client_tracker.lock().await.shutdown().await;
                return;
            }
        };
        self.pairing_persistence_key
            .send_replace(app_server_client_name.clone());
        if matches!(
            *self.desired_state_rx.borrow(),
            RemoteControlDesiredState::Unknown
        ) && !self
            .resolve_unknown_desired_state(app_server_client_name.as_deref())
            .await
        {
            self.client_tracker.lock().await.shutdown().await;
            return;
        }

        loop {
            if !self.wait_until_enabled().await {
                info!(
                    remote_control_url = %self.remote_control_url,
                    installation_id = %self.installation_id,
                    server_name = %self.server_name,
                    shutdown_requested = self.shutdown_token.is_cancelled(),
                    current_status = ?self.status_publisher.status().status,
                    "app-server remote control websocket loop exiting while waiting for enablement"
                );
                break;
            }

            let status = self.status_publisher.status();
            info!(
                remote_control_url = %self.remote_control_url,
                installation_id = %self.installation_id,
                server_name = %self.server_name,
                reconnect_attempt = self.reconnect_attempt.saturating_add(1),
                current_status = ?status.status,
                environment_id = ?status.environment_id,
                "starting app-server remote control websocket connection cycle"
            );
            let shutdown_token = self.shutdown_token.child_token();
            let websocket_connection = match self
                .connect(&shutdown_token, app_server_client_name.as_deref())
                .await
            {
                ConnectOutcome::Connected(websocket_connection) => *websocket_connection,
                ConnectOutcome::Disabled => {
                    self.status_publisher
                        .publish_status(RemoteControlConnectionStatus::Disabled);
                    continue;
                }
                ConnectOutcome::Shutdown => break,
            };

            let connection_end_reason = self
                .run_connection(websocket_connection, shutdown_token)
                .await;
            let status = self.status_publisher.status();
            info!(
                remote_control_url = %self.remote_control_url,
                installation_id = %self.installation_id,
                server_name = %self.server_name,
                connection_end_reason = ?connection_end_reason,
                current_status = ?status.status,
                environment_id = ?status.environment_id,
                desired_state = ?*self.desired_state_rx.borrow(),
                "app-server remote control websocket connection cycle ended"
            );
        }

        self.client_tracker.lock().await.shutdown().await;
        info!(
            remote_control_url = %self.remote_control_url,
            installation_id = %self.installation_id,
            server_name = %self.server_name,
            shutdown_requested = self.shutdown_token.is_cancelled(),
            "app-server remote control websocket loop exited"
        );
    }

    async fn wait_for_app_server_client_name(
        &self,
        app_server_client_name_rx: Option<oneshot::Receiver<String>>,
    ) -> Result<Option<String>, ()> {
        match app_server_client_name_rx {
            Some(app_server_client_name_rx) => {
                tokio::select! {
                    _ = self.shutdown_token.cancelled() => Err(()),
                    app_server_client_name = app_server_client_name_rx => match app_server_client_name {
                        Ok(app_server_client_name) => Ok(Some(app_server_client_name)),
                        Err(_) => Err(()),
                    },
                }
            }
            None => Ok(None),
        }
    }

    pub(super) async fn resolve_unknown_desired_state(
        &mut self,
        app_server_client_name: Option<&str>,
    ) -> bool {
        let remote_control_target = match super::protocol::normalize_remote_control_url(
            &self.remote_control_url,
        ) {
            Ok(remote_control_target) => remote_control_target,
            Err(err) => {
                warn!(
                    "remote control preference cannot be resolved because the URL is invalid: {err}"
                );
                self.transition_unknown_to(RemoteControlDesiredState::Disabled);
                return true;
            }
        };
        self.remote_control_target = Some(remote_control_target.clone());
        let Some(state_db) = self.state_db.clone() else {
            self.transition_unknown_to(RemoteControlDesiredState::Disabled);
            return true;
        };

        loop {
            if !matches!(
                *self.desired_state_rx.borrow(),
                RemoteControlDesiredState::Unknown
            ) {
                return true;
            }
            let auth = match load_remote_control_auth(&self.auth_manager).await {
                Ok(auth) => auth,
                Err(err) => {
                    info!(
                        error = %err,
                        "waiting to resolve remote control preference until authentication is available"
                    );
                    if !self.wait_for_preference_resolution_retry().await {
                        return false;
                    }
                    continue;
                }
            };
            let enrollment = match state_db
                .get_remote_control_enrollment(
                    &remote_control_target.websocket_url,
                    &auth.account_id,
                    app_server_client_name,
                )
                .await
            {
                Ok(enrollment) => enrollment,
                Err(err) => {
                    warn!(
                        error = %err,
                        "failed to resolve persisted remote control preference; retrying"
                    );
                    if !self.wait_for_preference_resolution_retry().await {
                        return false;
                    }
                    continue;
                }
            };
            let desired_state = desired_state_from_persisted_enrollment(enrollment);
            self.transition_unknown_to(desired_state);
            return true;
        }
    }

    fn transition_unknown_to(&self, desired_state: RemoteControlDesiredState) {
        self.desired_state_tx.send_if_modified(|state| {
            if !matches!(*state, RemoteControlDesiredState::Unknown) {
                return false;
            }
            *state = desired_state;
            true
        });
    }

    async fn wait_for_preference_resolution_retry(&mut self) -> bool {
        tokio::select! {
            _ = self.shutdown_token.cancelled() => false,
            changed = self.desired_state_rx.changed() => changed.is_ok(),
            _ = tokio::time::sleep(REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL) => true,
        }
    }

    async fn wait_until_enabled(&mut self) -> bool {
        tokio::select! {
            _ = self.shutdown_token.cancelled() => false,
            desired_state = self.desired_state_rx.wait_for(|state| state.is_enabled()) => desired_state.is_ok(),
        }
    }

    async fn connect(
        &mut self,
        shutdown_token: &CancellationToken,
        app_server_client_name: Option<&str>,
    ) -> ConnectOutcome {
        self.status_publisher
            .publish_status(RemoteControlConnectionStatus::Connecting);
        let remote_control_target = match self.remote_control_target.as_ref() {
            Some(remote_control_target) => remote_control_target.clone(),
            None => match super::protocol::normalize_remote_control_url(&self.remote_control_url) {
                Ok(remote_control_target) => {
                    self.remote_control_target = Some(remote_control_target.clone());
                    remote_control_target
                }
                Err(err) => {
                    self.status_publisher
                        .publish_status(RemoteControlConnectionStatus::Errored);
                    warn!("remote control is enabled but the URL is invalid: {err}");
                    tokio::select! {
                        _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown,
                        changed = self.desired_state_rx.wait_for(|state| !state.is_enabled()) => {
                            if changed.is_err() {
                                return ConnectOutcome::Shutdown;
                            }
                            return ConnectOutcome::Disabled;
                        }
                    }
                }
            },
        };

        loop {
            let subscribe_cursor = self.state.lock().await.subscribe_cursor.clone();
            let enrollment = self.current_enrollment.snapshot();
            info!(
                websocket_url = %remote_control_target.websocket_url,
                installation_id = %self.installation_id,
                server_name = %self.server_name,
                reconnect_attempt = self.reconnect_attempt.saturating_add(1),
                has_enrollment = enrollment.is_some(),
                server_id = ?enrollment.as_ref().map(|enrollment| enrollment.server_id.as_str()),
                environment_id = ?enrollment.as_ref().map(|enrollment| enrollment.environment_id.as_str()),
                subscribe_cursor_present = subscribe_cursor.is_some(),
                app_server_client_name = ?app_server_client_name,
                "connecting to app-server remote control websocket"
            );
            let connect_options = RemoteControlConnectOptions {
                installation_id: &self.installation_id,
                server_name: &self.server_name,
                subscribe_cursor: subscribe_cursor.as_deref(),
                app_server_client_name,
                desired_state_tx: &self.desired_state_tx,
                desired_state_persistence_lock: &self.desired_state_persistence_lock,
            };
            let auth_context = RemoteControlAuthContext {
                auth_manager: &self.auth_manager,
                auth_recovery: &mut self.auth_recovery,
                auth_change_rx: &mut self.auth_change_rx,
            };
            let mut disabled_rx = self.desired_state_rx.clone();
            let connect_result = tokio::select! {
                _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown,
                changed = disabled_rx.wait_for(|state| !state.is_enabled()) => {
                    if changed.is_err() {
                        return ConnectOutcome::Shutdown;
                    }
                    return ConnectOutcome::Disabled;
                }
                connect_result = async {
                    connect_remote_control_websocket(
                        &remote_control_target,
                        self.state_db.as_deref(),
                        auth_context,
                        &self.current_enrollment,
                        connect_options,
                        &self.status_publisher,
                    )
                    .await
                } => connect_result,
            };

            match connect_result {
                Ok((websocket_connection, response)) => {
                    if !self.desired_state_rx.borrow().is_enabled() {
                        return ConnectOutcome::Disabled;
                    }
                    self.reconnect_attempt = 0;
                    self.auth_recovery = self.auth_manager.unauthorized_recovery();
                    self.status_publisher
                        .publish_status(RemoteControlConnectionStatus::Connected);
                    let enrollment = self.current_enrollment.snapshot();
                    info!(
                        websocket_url = %remote_control_target.websocket_url,
                        installation_id = %self.installation_id,
                        server_name = %self.server_name,
                        server_id = ?enrollment.as_ref().map(|enrollment| enrollment.server_id.as_str()),
                        environment_id = ?enrollment.as_ref().map(|enrollment| enrollment.environment_id.as_str()),
                        subscribe_cursor_present = subscribe_cursor.is_some(),
                        response_headers = %format_headers(response.headers()),
                        "connected to app-server remote control websocket"
                    );
                    return ConnectOutcome::Connected(Box::new(websocket_connection));
                }
                Err(err) => {
                    if !self.desired_state_rx.borrow().is_enabled() {
                        return ConnectOutcome::Disabled;
                    }
                    let reconnect_delay = if err.kind() == ErrorKind::WouldBlock {
                        REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL
                    } else {
                        self.status_publisher
                            .publish_status(RemoteControlConnectionStatus::Errored);
                        let reconnect_attempt = self.reconnect_attempt.saturating_add(1);
                        let (reconnect_delay, reconnect_backoff_reset) =
                            next_reconnect_delay(&mut self.reconnect_attempt);
                        let enrollment = self.current_enrollment.snapshot();
                        warn!(
                            websocket_url = %remote_control_target.websocket_url,
                            installation_id = %self.installation_id,
                            server_name = %self.server_name,
                            error = %err,
                            error_kind = ?err.kind(),
                            reconnect_attempt,
                            reconnect_delay = ?reconnect_delay,
                            reconnect_backoff_reset,
                            has_enrollment = enrollment.is_some(),
                            server_id = ?enrollment.as_ref().map(|enrollment| enrollment.server_id.as_str()),
                            environment_id = ?enrollment.as_ref().map(|enrollment| enrollment.environment_id.as_str()),
                            subscribe_cursor_present = subscribe_cursor.is_some(),
                            "failed to connect to app-server remote control websocket"
                        );
                        if reconnect_backoff_reset {
                            info!(
                                reconnect_backoff_cap = ?REMOTE_CONTROL_RECONNECT_BACKOFF_CAP,
                                "reset app-server remote control websocket reconnect backoff after cap"
                            );
                        }
                        reconnect_delay
                    };
                    tokio::select! {
                        _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown,
                        changed = self.desired_state_rx.wait_for(|state| !state.is_enabled()) => {
                            if changed.is_err() {
                                return ConnectOutcome::Shutdown;
                            }
                            return ConnectOutcome::Disabled;
                        }
                        changed = self.auth_change_rx.changed() => {
                            if changed.is_err() {
                                return ConnectOutcome::Shutdown;
                            }
                            self.auth_recovery = self.auth_manager.unauthorized_recovery();
                            self.reconnect_attempt = 0;
                            info!("retrying app-server remote control websocket after auth changed");
                        }
                        _ = tokio::time::sleep(reconnect_delay) => {}
                    }
                }
            }
        }
    }

    async fn run_connection(
        &self,
        websocket_connection: WebSocketStream<MaybeTlsStream<TcpStream>>,
        shutdown_token: CancellationToken,
    ) -> ConnectionEndReason {
        let (websocket_writer, websocket_reader) = websocket_connection.split();
        let mut join_set = tokio::task::JoinSet::new();

        join_set.spawn(Self::run_server_writer(
            self.state.clone(),
            self.server_event_rx.clone(),
            self.used_rx.clone(),
            websocket_writer,
            REMOTE_CONTROL_WEBSOCKET_PING_INTERVAL,
            shutdown_token.clone(),
        ));
        join_set.spawn(Self::run_websocket_reader(
            self.client_tracker.clone(),
            self.state.clone(),
            websocket_reader,
            REMOTE_CONTROL_WEBSOCKET_PONG_TIMEOUT,
            shutdown_token.clone(),
        ));

        let mut desired_state_rx = self.desired_state_rx.clone();
        let connection_end_reason = tokio::select! {
            _ = shutdown_token.cancelled() => ConnectionEndReason::Shutdown,
            changed = desired_state_rx.wait_for(|state| !state.is_enabled()) => {
                if changed.is_ok() {
                    self.status_publisher
                        .publish_status(RemoteControlConnectionStatus::Disabled);
                    ConnectionEndReason::Disabled
                } else {
                    ConnectionEndReason::EnabledWatchClosed
                }
            }
            _ = join_set.join_next() => ConnectionEndReason::ConnectionWorkerStopped,
        };
        shutdown_token.cancel();

        Self::join_connection_workers(&mut join_set, REMOTE_CONTROL_CONNECTION_SHUTDOWN_TIMEOUT)
            .await;
        connection_end_reason
    }

    async fn join_connection_workers(
        join_set: &mut tokio::task::JoinSet<()>,
        shutdown_timeout: std::time::Duration,
    ) {
        if tokio::time::timeout(shutdown_timeout, Self::drain_join_set(join_set))
            .await
            .is_ok()
        {
            return;
        }

        warn!(
            shutdown_timeout = ?shutdown_timeout,
            remaining_workers = join_set.len(),
            "timed out waiting for remote control connection workers to stop; aborting"
        );
        join_set.abort_all();
        Self::drain_join_set(join_set).await;
    }

    async fn drain_join_set(join_set: &mut tokio::task::JoinSet<()>) {
        while join_set.join_next().await.is_some() {}
    }

    async fn run_server_writer(
        state: Arc<Mutex<WebsocketState>>,
        server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>,
        used_rx: watch::Receiver<usize>,
        websocket_writer: SplitSink<
            WebSocketStream<MaybeTlsStream<TcpStream>>,
            tungstenite::Message,
        >,
        ping_interval: std::time::Duration,
        shutdown_token: CancellationToken,
    ) {
        let result = Self::run_server_writer_inner(
            state,
            server_event_rx,
            used_rx,
            websocket_writer,
            ping_interval,
            shutdown_token,
        )
        .await;
        if let Err(err) = result {
            warn!("remote control websocket writer disconnected, err: {err}");
        } else {
            warn!("remote control websocket writer was stopped");
        }
    }

    #[expect(
        clippy::await_holding_invalid_type,
        reason = "remote-control server event receiver is shared across reconnects"
    )]
    async fn run_server_writer_inner(
        state: Arc<Mutex<WebsocketState>>,
        server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>,
        mut used_rx: watch::Receiver<usize>,
        mut websocket_writer: SplitSink<
            WebSocketStream<MaybeTlsStream<TcpStream>>,
            tungstenite::Message,
        >,
        ping_interval: std::time::Duration,
        shutdown_token: CancellationToken,
    ) -> io::Result<()> {
        let server_envelopes = state
            .lock()
            .await
            .outbound_buffer
            .server_envelopes()
            .cloned()
            .collect::<Vec<_>>();
        for server_envelope in server_envelopes {
            let payload = match serde_json::to_string(&server_envelope) {
                Ok(payload) => payload,
                Err(err) => {
                    error!("failed to serialize remote-control server event: {err}");
                    continue;
                }
            };
            tokio::select! {
                _ = shutdown_token.cancelled() => return Ok(()),
                send_result = websocket_writer.send(tungstenite::Message::Text(payload.into())) => {
                    if let Err(err) = send_result {
                        return Err(io::Error::other(err));
                    }
                }
            };
        }

        let mut ping_interval =
            tokio::time::interval_at(tokio::time::Instant::now() + ping_interval, ping_interval);
        ping_interval.set_missed_tick_behavior(MissedTickBehavior::Skip);

        let mut server_event_rx = server_event_rx.lock().await;
        loop {
            let outbound_has_capacity = *used_rx.borrow() < super::CHANNEL_CAPACITY;
            let queued_server_envelope = tokio::select! {
                _ = shutdown_token.cancelled() => return Ok(()),
                _ = ping_interval.tick() => {
                    tokio::select! {
                        _ = shutdown_token.cancelled() => return Ok(()),
                        send_result = websocket_writer.send(tungstenite::Message::Ping(Vec::new().into())) => {
                            if let Err(err) = send_result {
                                return Err(io::Error::other(err));
                            }
                        }
                    };
                    continue;
                }
                wait_result = used_rx.changed(), if !outbound_has_capacity =>
                {
                    if wait_result.is_err() {
                        return Err(io::Error::new(
                            ErrorKind::UnexpectedEof,
                            "outbound buffer usage channel closed",
                        ));
                    }
                    continue;
                }
                recv_result = server_event_rx.recv(), if outbound_has_capacity => {
                    match recv_result {
                        Some(queued_server_envelope) => queued_server_envelope,
                        None => {
                            return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "server event channel closed"));
                        }
                    }
                }
            };
            let (payloads, write_complete_tx) = {
                let mut state = state.lock().await;
                let seq_key = (
                    queued_server_envelope.client_id.clone(),
                    queued_server_envelope.stream_id.clone(),
                );
                let seq_id = *state
                    .next_seq_id_by_stream
                    .entry(seq_key.clone())
                    .or_insert(1);

                let server_envelope = ServerEnvelope {
                    event: queued_server_envelope.event,
                    client_id: queued_server_envelope.client_id,
                    seq_id,
                    stream_id: queued_server_envelope.stream_id,
                };
                let server_envelopes = match split_server_envelope_for_transport(server_envelope) {
                    Ok(server_envelopes) => server_envelopes,
                    Err(err) => {
                        error!("failed to split remote-control server event: {err}");
                        continue;
                    }
                };
                let mut payloads = Vec::with_capacity(server_envelopes.len());
                for server_envelope in server_envelopes {
                    let payload = match serde_json::to_string(&server_envelope) {
                        Ok(payload) => payload,
                        Err(err) => {
                            error!("failed to serialize remote-control server event: {err}");
                            continue;
                        }
                    };
                    state.outbound_buffer.insert(&server_envelope);
                    payloads.push(payload);
                }
                state
                    .next_seq_id_by_stream
                    .insert(seq_key, seq_id.saturating_add(1));

                (payloads, queued_server_envelope.write_complete_tx)
            };

            for payload in payloads {
                tokio::select! {
                    _ = shutdown_token.cancelled() => return Ok(()),
                    send_result = websocket_writer.send(tungstenite::Message::Text(payload.into())) => {
                        if let Err(err) = send_result {
                            return Err(io::Error::other(err));
                        }
                    }
                }
            }
            if let Some(write_complete_tx) = write_complete_tx {
                let _ = write_complete_tx.send(());
            }
        }
    }

    async fn run_websocket_reader(
        client_tracker: Arc<Mutex<ClientTracker>>,
        state: Arc<Mutex<WebsocketState>>,
        websocket_reader: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
        pong_timeout: std::time::Duration,
        shutdown_token: CancellationToken,
    ) {
        let result = Self::run_websocket_reader_inner(
            client_tracker,
            state,
            websocket_reader,
            pong_timeout,
            shutdown_token,
        )
        .await;
        if let Err(err) = result {
            warn!("remote control websocket reader disconnected, err: {err}");
        } else {
            warn!("remote control websocket reader was stopped");
        }
    }

    #[expect(
        clippy::await_holding_invalid_type,
        reason = "remote-control client tracking must stay serialized while processing inbound events"
    )]
    async fn run_websocket_reader_inner(
        client_tracker: Arc<Mutex<ClientTracker>>,
        state: Arc<Mutex<WebsocketState>>,
        mut websocket_reader: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>,
        pong_timeout: std::time::Duration,
        shutdown_token: CancellationToken,
    ) -> io::Result<()> {
        let mut client_tracker = client_tracker.lock().await;
        let mut idle_sweep_interval = tokio::time::interval(REMOTE_CONTROL_IDLE_SWEEP_INTERVAL);
        idle_sweep_interval.set_missed_tick_behavior(MissedTickBehavior::Skip);
        let pong_deadline = tokio::time::sleep(pong_timeout);
        tokio::pin!(pong_deadline);

        loop {
            let incoming_message = tokio::select! {
                _ = shutdown_token.cancelled() => return Ok(()),
                _ = &mut pong_deadline => {
                    return Err(io::Error::new(
                        ErrorKind::TimedOut,
                        "remote control websocket pong timeout",
                    ));
                }
                client_key = client_tracker.bookkeep_join_set() => {
                    let Some(client_key) = client_key else {
                        continue;
                    };
                    if client_tracker.close_client(&client_key).await.is_err() {
                        return Ok(());
                    }
                    state
                        .lock()
                        .await
                        .client_segment_reassembler
                        .invalidate_stream(&client_key.0, &client_key.1);
                    state
                        .lock()
                        .await
                        .invalidate_client_message_stream(&client_key.0, &client_key.1);
                    continue;
                }
                _ = idle_sweep_interval.tick() => {
                    match client_tracker.close_expired_clients().await {
                        Ok(client_keys) => {
                            let mut websocket_state = state.lock().await;
                            for (client_id, stream_id) in client_keys {
                                websocket_state
                                    .client_segment_reassembler
                                    .invalidate_stream(&client_id, &stream_id);
                                websocket_state
                                    .invalidate_client_message_stream(&client_id, &stream_id);
                            }
                        }
                        Err(_) => return Ok(()),
                    }
                    continue;
                }
                incoming_message = websocket_reader.next() => {
                    match incoming_message {
                        Some(incoming_message) => incoming_message,
                        None => return Err(io::Error::new(ErrorKind::UnexpectedEof, "websocket stream ended")),
                    }
                }
            };
            let (client_envelope, wire_size_bytes) = match incoming_message {
                Ok(tungstenite::Message::Text(text)) => {
                    let wire_size_bytes = text.len();
                    match serde_json::from_str::<ClientEnvelope>(&text) {
                        Ok(client_envelope) => (client_envelope, wire_size_bytes),
                        Err(err) => {
                            warn!("failed to deserialize remote-control client event: {err}");
                            continue;
                        }
                    }
                }
                Ok(tungstenite::Message::Pong(_)) => {
                    pong_deadline
                        .as_mut()
                        .reset(tokio::time::Instant::now() + pong_timeout);
                    continue;
                }
                Ok(tungstenite::Message::Ping(_)) | Ok(tungstenite::Message::Frame(_)) => continue,
                Ok(tungstenite::Message::Binary(_)) => {
                    warn!("dropping unsupported binary remote-control websocket message");
                    continue;
                }
                Ok(tungstenite::Message::Close(_)) => {
                    return Err(io::Error::new(
                        ErrorKind::ConnectionAborted,
                        "websocket disconnected",
                    ));
                }
                Err(err) => {
                    return Err(io::Error::new(
                        ErrorKind::InvalidData,
                        format!("failed to read from websocket: {err}"),
                    ));
                }
            };

            let client_message_key = WebsocketState::client_message_key(&client_envelope);
            let observation = {
                let mut websocket_state = state.lock().await;
                websocket_state.observe_client_message(client_envelope, wire_size_bytes)
            };
            let client_envelope = match observation {
                ClientSegmentObservation::Forward(client_envelope) => *client_envelope,
                ClientSegmentObservation::Pending | ClientSegmentObservation::Dropped => continue,
            };

            let closed_client =
                matches!(&client_envelope.event, ClientEvent::ClientClosed).then(|| {
                    (
                        client_envelope.client_id.clone(),
                        client_envelope.stream_id.clone(),
                    )
                });
            let delivered_client_envelope = client_envelope.clone();
            if client_tracker
                .handle_message(client_envelope)
                .await
                .is_err()
            {
                return Ok(());
            }
            state
                .lock()
                .await
                .record_client_message_delivery(&delivered_client_envelope, client_message_key);
            if let Some((client_id, stream_id)) = closed_client {
                let mut websocket_state = state.lock().await;
                if let Some(stream_id) = stream_id {
                    websocket_state
                        .client_segment_reassembler
                        .invalidate_stream(&client_id, &stream_id);
                    websocket_state.invalidate_client_message_stream(&client_id, &stream_id);
                } else {
                    websocket_state
                        .client_segment_reassembler
                        .invalidate_client(&client_id);
                    websocket_state.invalidate_client_message_client(&client_id);
                }
            }
        }
    }
}

fn set_remote_control_header(
    headers: &mut tungstenite::http::HeaderMap,
    name: &'static str,
    value: &str,
) -> io::Result<()> {
    let header_value = HeaderValue::from_str(value).map_err(|err| {
        io::Error::new(
            ErrorKind::InvalidInput,
            format!("invalid remote control header `{name}`: {err}"),
        )
    })?;
    headers.insert(name, header_value);
    Ok(())
}

fn build_remote_control_websocket_request(
    websocket_url: &str,
    enrollment: &RemoteControlEnrollment,
    installation_id: &str,
    subscribe_cursor: Option<&str>,
) -> io::Result<tungstenite::http::Request<()>> {
    let mut request = websocket_url.into_client_request().map_err(|err| {
        io::Error::new(
            ErrorKind::InvalidInput,
            format!("invalid remote control websocket URL `{websocket_url}`: {err}"),
        )
    })?;
    let headers = request.headers_mut();
    set_remote_control_header(headers, "x-codex-server-id", &enrollment.server_id)?;
    set_remote_control_header(
        headers,
        "x-codex-name",
        &base64::engine::general_purpose::STANDARD.encode(&enrollment.server_name),
    )?;
    set_remote_control_header(
        headers,
        "x-codex-protocol-version",
        REMOTE_CONTROL_PROTOCOL_VERSION,
    )?;
    set_remote_control_header(
        headers,
        "authorization",
        &format!(
            "Bearer {}",
            enrollment
                .remote_control_token
                .as_deref()
                .ok_or_else(|| io::Error::other("missing remote control server token"))?
        ),
    )?;
    set_remote_control_header(
        headers,
        REMOTE_CONTROL_INSTALLATION_ID_HEADER,
        installation_id,
    )?;
    if let Some(subscribe_cursor) = subscribe_cursor {
        set_remote_control_header(
            headers,
            REMOTE_CONTROL_SUBSCRIBE_CURSOR_HEADER,
            subscribe_cursor,
        )?;
    }
    Ok(request)
}

fn next_reconnect_delay(reconnect_attempt: &mut u64) -> (std::time::Duration, bool) {
    let reconnect_delay = backoff(*reconnect_attempt).min(REMOTE_CONTROL_RECONNECT_BACKOFF_CAP);
    let reconnect_backoff_reset = reconnect_delay == REMOTE_CONTROL_RECONNECT_BACKOFF_CAP;
    *reconnect_attempt = if reconnect_backoff_reset {
        0
    } else {
        (*reconnect_attempt).saturating_add(1)
    };
    (reconnect_delay, reconnect_backoff_reset)
}

pub(super) async fn connect_remote_control_websocket(
    remote_control_target: &RemoteControlTarget,
    state_db: Option<&StateRuntime>,
    mut auth_context: RemoteControlAuthContext<'_>,
    current_enrollment: &CurrentRemoteControlEnrollment,
    connect_options: RemoteControlConnectOptions<'_>,
    status_publisher: &RemoteControlStatusPublisher,
) -> io::Result<(
    WebSocketStream<MaybeTlsStream<TcpStream>>,
    tungstenite::http::Response<()>,
)> {
    ensure_rustls_crypto_provider();

    let (auth, enrollment) = {
        let mut current_enrollment = current_enrollment.lock().await;
        let auth = prepare_remote_control_enrollment(
            remote_control_target,
            state_db,
            &mut auth_context,
            &mut current_enrollment,
            connect_options,
            status_publisher,
        )
        .await?;
        let enrollment = current_enrollment.as_ref().cloned().ok_or_else(|| {
            io::Error::other("missing remote control enrollment after enrollment step")
        })?;
        (auth, enrollment)
    };
    let request = build_remote_control_websocket_request(
        &remote_control_target.websocket_url,
        &enrollment,
        connect_options.installation_id,
        connect_options.subscribe_cursor,
    )?;

    let websocket_connect_result = tokio::time::timeout(
        REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT,
        connect_async(request),
    )
    .await
    .map_err(|_| {
        io::Error::new(
            ErrorKind::TimedOut,
            format!(
                "timed out connecting to remote control websocket at `{}` after {:?}",
                remote_control_target.websocket_url, REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT
            ),
        )
    })?;

    match websocket_connect_result {
        Ok((websocket_stream, response)) => Ok((websocket_stream, response.map(|_| ()))),
        Err(err) => {
            match &err {
                tungstenite::Error::Http(response)
                    if websocket_response_reports_missing_remote_app_server(response) =>
                {
                    info!(
                        "remote control websocket returned HTTP 404; replacing stale enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}",
                        remote_control_target.websocket_url,
                        auth.account_id,
                        enrollment.server_id,
                        enrollment.environment_id
                    );
                    replace_remote_control_enrollment_if_matches(
                        state_db,
                        remote_control_target,
                        RemoteControlEnrollmentAuthContext {
                            auth: &auth,
                            recovery: &mut auth_context,
                        },
                        current_enrollment,
                        &enrollment,
                        connect_options,
                        status_publisher,
                    )
                    .await?;
                }
                tungstenite::Error::Http(response) if response.status().as_u16() == 404 => {
                    let response_body = response
                        .body()
                        .as_deref()
                        .map(preview_remote_control_response_body)
                        .unwrap_or_else(|| "<missing>".to_string());
                    warn!(
                        websocket_url = %remote_control_target.websocket_url,
                        account_id = %auth.account_id,
                        server_id = %enrollment.server_id,
                        environment_id = %enrollment.environment_id,
                        response_status = %response.status(),
                        response_headers = %format_headers(response.headers()),
                        response_body = %response_body,
                        "remote control websocket returned unrecognized HTTP 404; preserving enrollment before retry"
                    );
                }
                tungstenite::Error::Http(response)
                    if matches!(response.status().as_u16(), 401 | 403) =>
                {
                    clear_remote_control_server_token_if_matches(current_enrollment, &enrollment)
                        .await?;
                    return Err(io::Error::other(format!(
                        "remote control websocket auth failed with HTTP {}; refreshing server token before reconnect",
                        response.status()
                    )));
                }
                _ => {}
            }
            Err(io::Error::other(
                format_remote_control_websocket_connect_error(
                    &remote_control_target.websocket_url,
                    &err,
                ),
            ))
        }
    }
}

async fn prepare_remote_control_enrollment(
    remote_control_target: &RemoteControlTarget,
    state_db: Option<&StateRuntime>,
    auth_context: &mut RemoteControlAuthContext<'_>,
    enrollment: &mut Option<RemoteControlEnrollment>,
    connect_options: RemoteControlConnectOptions<'_>,
    status_publisher: &RemoteControlStatusPublisher,
) -> io::Result<RemoteControlConnectionAuth> {
    let Some(state_db) = state_db else {
        *enrollment = None;
        return Err(io::Error::new(
            ErrorKind::NotFound,
            "remote control requires sqlite state db",
        ));
    };

    let auth = match load_remote_control_auth(auth_context.auth_manager).await {
        Ok(auth) => auth,
        Err(err) => {
            if err.kind() == ErrorKind::PermissionDenied {
                *enrollment = None;
                status_publisher.publish_environment_id(/*environment_id*/ None);
            }
            return Err(err);
        }
    };
    let enrollment_account_id = enrollment.as_ref().map(|enrollment| &enrollment.account_id);
    if enrollment_account_id.is_some_and(|account_id| account_id != &auth.account_id) {
        resolve_desired_state_after_account_change(
            state_db,
            remote_control_target,
            auth_context.auth_manager,
            &auth.account_id,
            connect_options,
        )
        .await?;
        info!(
            "clearing in-memory remote control enrollment because account id changed: websocket_url={}, previous_account_id={:?}, current_account_id={:?}",
            remote_control_target.websocket_url,
            enrollment
                .as_ref()
                .map(|enrollment| enrollment.account_id.as_str()),
            auth.account_id
        );
        *enrollment = None;
        status_publisher.publish_environment_id(/*environment_id*/ None);
        if !connect_options.desired_state_tx.borrow().is_enabled() {
            return Err(io::Error::new(
                ErrorKind::Interrupted,
                "remote control disabled after account changed",
            ));
        }
    }
    if let Some(enrollment) = enrollment.as_mut() {
        enrollment.remote_control_target = remote_control_target.clone();
    }

    if let Some(enrollment) = enrollment.as_ref() {
        status_publisher.publish_environment_id(Some(enrollment.environment_id.clone()));
    }

    if enrollment.is_none() {
        let loaded_enrollment = load_persisted_remote_control_enrollment(
            Some(state_db),
            remote_control_target,
            &auth.account_id,
            connect_options.app_server_client_name,
        )
        .await?;
        if let Some(loaded_enrollment) = loaded_enrollment.as_ref() {
            status_publisher.publish_environment_id(Some(loaded_enrollment.environment_id.clone()));
        }
        *enrollment = loaded_enrollment.map(|mut enrollment| {
            enrollment.server_name = connect_options.server_name.to_string();
            enrollment
        });
    }

    enroll_and_persist_remote_control_server(
        remote_control_target,
        state_db,
        RemoteControlEnrollmentAuthContext {
            auth: &auth,
            recovery: auth_context,
        },
        enrollment,
        connect_options,
        status_publisher,
        RemoteControlEnrollmentSelection::ReuseOrCreate,
    )
    .await?;

    if enrollment
        .as_ref()
        .ok_or_else(|| io::Error::other("missing remote control enrollment after enrollment step"))?
        .should_refresh_server_token()
    {
        let enrollment_ref = enrollment.as_ref().ok_or_else(|| {
            io::Error::other("missing remote control enrollment after enrollment step")
        })?;
        let server_id = enrollment_ref.server_id.clone();
        let environment_id = enrollment_ref.environment_id.clone();

        info!(
            "refreshing remote control server token: websocket_url={}, refresh_url={}, account_id={}, server_id={}, environment_id={}",
            remote_control_target.websocket_url,
            remote_control_target.refresh_url,
            auth.account_id,
            server_id,
            environment_id
        );
        let enrollment_ref = enrollment.as_mut().ok_or_else(|| {
            io::Error::other("missing remote control enrollment before server refresh")
        })?;
        match refresh_remote_control_server(&auth, connect_options.installation_id, enrollment_ref)
            .await
        {
            Ok(()) => {}
            Err(err) if err.kind() == ErrorKind::NotFound => {
                info!(
                    "remote control server refresh returned HTTP 404; replacing stale enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}",
                    remote_control_target.websocket_url, auth.account_id, server_id, environment_id
                );
                enroll_and_persist_remote_control_server(
                    remote_control_target,
                    state_db,
                    RemoteControlEnrollmentAuthContext {
                        auth: &auth,
                        recovery: auth_context,
                    },
                    enrollment,
                    connect_options,
                    status_publisher,
                    RemoteControlEnrollmentSelection::ReplaceExisting,
                )
                .await?;
            }
            Err(err)
                if err.kind() == ErrorKind::PermissionDenied
                    && recover_remote_control_auth(
                        auth_context.auth_recovery,
                        auth_context.auth_change_rx,
                    )
                    .await =>
            {
                return Err(io::Error::other(format!(
                    "{err}; retrying after auth recovery"
                )));
            }
            Err(err) => return Err(err),
        }
    }

    Ok(auth)
}

async fn resolve_desired_state_after_account_change(
    state_db: &StateRuntime,
    remote_control_target: &RemoteControlTarget,
    auth_manager: &Arc<AuthManager>,
    account_id: &str,
    connect_options: RemoteControlConnectOptions<'_>,
) -> io::Result<()> {
    let durable_enabled = RemoteControlDesiredState::Enabled {
        persistence_preference: Some(true),
    };
    if *connect_options.desired_state_tx.borrow() != durable_enabled {
        return Ok(());
    }

    let _persistence =
        acquire_persistence_lock(connect_options.desired_state_persistence_lock).await;
    if *connect_options.desired_state_tx.borrow() != durable_enabled {
        return Ok(());
    }
    let enrollment = state_db
        .get_remote_control_enrollment(
            &remote_control_target.websocket_url,
            account_id,
            connect_options.app_server_client_name,
        )
        .await
        .map_err(io::Error::other)?;
    let current_auth = load_remote_control_auth(auth_manager).await?;
    if current_auth.account_id != account_id {
        return Err(io::Error::new(
            ErrorKind::WouldBlock,
            "remote control account changed while resolving persisted preference",
        ));
    }
    let resolved_state = desired_state_from_persisted_enrollment(enrollment);
    connect_options.desired_state_tx.send_if_modified(|state| {
        if *state != durable_enabled || *state == resolved_state {
            return false;
        }
        *state = resolved_state;
        true
    });
    Ok(())
}

fn websocket_response_reports_missing_remote_app_server(
    response: &tungstenite::http::Response<Option<Vec<u8>>>,
) -> bool {
    response.status().as_u16() == 404
        && response.body().as_deref().is_some_and(|body| {
            serde_json::from_slice::<serde_json::Value>(body).is_ok_and(|body| {
                body.get("detail").and_then(serde_json::Value::as_str)
                    == Some(REMOTE_APP_SERVER_NOT_FOUND_DETAIL)
            })
        })
}

async fn replace_remote_control_enrollment_if_matches(
    state_db: Option<&StateRuntime>,
    remote_control_target: &RemoteControlTarget,
    auth_context: RemoteControlEnrollmentAuthContext<'_, '_>,
    current_enrollment: &CurrentRemoteControlEnrollment,
    enrollment: &RemoteControlEnrollment,
    connect_options: RemoteControlConnectOptions<'_>,
    status_publisher: &RemoteControlStatusPublisher,
) -> io::Result<()> {
    let Some(state_db) = state_db else {
        return Err(io::Error::new(
            ErrorKind::NotFound,
            "remote control requires sqlite state db",
        ));
    };
    let mut current_enrollment = current_enrollment.lock().await;
    if !current_enrollment
        .as_ref()
        .is_some_and(|current| same_remote_control_enrollment(current, enrollment))
    {
        return Ok(());
    }
    enroll_and_persist_remote_control_server(
        remote_control_target,
        state_db,
        auth_context,
        &mut current_enrollment,
        connect_options,
        status_publisher,
        RemoteControlEnrollmentSelection::ReplaceExisting,
    )
    .await
}

async fn clear_remote_control_server_token_if_matches(
    current_enrollment: &CurrentRemoteControlEnrollment,
    enrollment: &RemoteControlEnrollment,
) -> io::Result<()> {
    let mut current_enrollment = current_enrollment.lock().await;
    current_enrollment
        .as_mut()
        .filter(|current| same_remote_control_enrollment(current, enrollment))
        .ok_or_else(|| {
            io::Error::other("missing remote control enrollment after websocket auth failure")
        })?
        .clear_server_token();
    Ok(())
}

async fn enroll_and_persist_remote_control_server(
    remote_control_target: &RemoteControlTarget,
    state_db: &StateRuntime,
    auth_context: RemoteControlEnrollmentAuthContext<'_, '_>,
    enrollment: &mut Option<RemoteControlEnrollment>,
    connect_options: RemoteControlConnectOptions<'_>,
    status_publisher: &RemoteControlStatusPublisher,
    selection: RemoteControlEnrollmentSelection,
) -> io::Result<()> {
    match selection {
        RemoteControlEnrollmentSelection::ReuseOrCreate => {
            if enrollment.is_some() {
                return Ok(());
            }
        }
        RemoteControlEnrollmentSelection::ReplaceExisting => {}
    }
    if !connect_options.desired_state_tx.borrow().is_enabled() {
        return Err(io::Error::new(
            ErrorKind::Interrupted,
            "remote control disabled before enrollment",
        ));
    }

    info!(
        "creating new remote control enrollment: websocket_url={}, enroll_url={}, account_id={}",
        remote_control_target.websocket_url,
        remote_control_target.enroll_url,
        auth_context.auth.account_id
    );
    let new_enrollment = match enroll_remote_control_server(
        remote_control_target,
        auth_context.auth,
        connect_options.installation_id,
        connect_options.server_name,
    )
    .await
    {
        Ok(new_enrollment) => new_enrollment,
        Err(err)
            if err.kind() == ErrorKind::PermissionDenied
                && recover_remote_control_auth(
                    auth_context.recovery.auth_recovery,
                    auth_context.recovery.auth_change_rx,
                )
                .await =>
        {
            return Err(io::Error::other(format!(
                "{err}; retrying after auth recovery"
            )));
        }
        Err(err) => return Err(err),
    };
    let _persistence =
        acquire_persistence_lock(connect_options.desired_state_persistence_lock).await;
    let persistence_preference = match *connect_options.desired_state_tx.borrow() {
        RemoteControlDesiredState::Enabled {
            persistence_preference,
        } => persistence_preference,
        RemoteControlDesiredState::Unknown | RemoteControlDesiredState::Disabled => {
            return Err(io::Error::new(
                ErrorKind::Interrupted,
                "remote control disabled during enrollment",
            ));
        }
    };
    if let Err(err) = update_persisted_remote_control_enrollment(
        Some(state_db),
        remote_control_target,
        &auth_context.auth.account_id,
        connect_options.app_server_client_name,
        Some(&new_enrollment),
        persistence_preference,
    )
    .await
    {
        return Err(io::Error::other(format!(
            "failed to persist remote control enrollment in sqlite state db: {err}"
        )));
    }
    info!(
        "created new remote control enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}",
        remote_control_target.websocket_url,
        new_enrollment.account_id,
        new_enrollment.server_id,
        new_enrollment.environment_id
    );
    status_publisher.publish_environment_id(Some(new_enrollment.environment_id.clone()));
    *enrollment = Some(new_enrollment);
    Ok(())
}

fn format_remote_control_websocket_connect_error(
    websocket_url: &str,
    err: &tungstenite::Error,
) -> String {
    let mut message =
        format!("failed to connect app-server remote control websocket `{websocket_url}`: {err}");
    let tungstenite::Error::Http(response) = err else {
        return message;
    };

    message.push_str(&format!(", {}", format_headers(response.headers())));
    if let Some(body) = response.body().as_ref()
        && !body.is_empty()
    {
        let body_preview = preview_remote_control_response_body(body);
        message.push_str(&format!(", body: {body_preview}"));
    }

    message
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::outgoing_message::OutgoingMessage;
    use crate::transport::remote_control::ServerEvent;
    use crate::transport::remote_control::auth::mark_recovery_auth_change_seen;
    use crate::transport::remote_control::protocol::StreamId;
    use crate::transport::remote_control::protocol::normalize_remote_control_url;
    use chrono::Utc;
    use codex_app_server_protocol::AuthMode;
    use codex_app_server_protocol::ConfigWarningNotification;
    use codex_app_server_protocol::JSONRPCMessage;
    use codex_app_server_protocol::JSONRPCNotification;
    use codex_app_server_protocol::ServerNotification;
    use codex_config::types::AuthCredentialsStoreMode;
    use codex_core::test_support::auth_manager_from_auth;
    use codex_login::AuthDotJson;
    use codex_login::AuthKeyringBackendKind;
    use codex_login::CodexAuth;
    use codex_login::save_auth;
    use codex_login::token_data::TokenData;
    use codex_login::token_data::parse_chatgpt_jwt_claims;
    use codex_state::StateRuntime;
    use futures::StreamExt;
    use pretty_assertions::assert_eq;
    use std::sync::Arc;
    use tempfile::TempDir;
    use tokio::io::AsyncBufReadExt;
    use tokio::io::AsyncWriteExt;
    use tokio::io::BufReader;
    use tokio::net::TcpListener;
    use tokio::net::TcpStream;
    use tokio::sync::mpsc;
    use tokio::time::Duration;
    use tokio::time::timeout;
    use tokio_tungstenite::accept_async;

    // Windows Bazel CI can take longer than a few seconds for the websocket
    // client connection attempt to reach the local test listener.
    #[cfg(windows)]
    const TEST_HTTP_ACCEPT_TIMEOUT: Duration = Duration::from_secs(30);
    #[cfg(not(windows))]
    const TEST_HTTP_ACCEPT_TIMEOUT: Duration = Duration::from_secs(5);
    const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111";
    const TEST_REMOTE_CONTROL_SERVER_TOKEN: &str = "Remote Control Token";

    fn remote_control_enrollment(remote_control_token: Option<&str>) -> RemoteControlEnrollment {
        RemoteControlEnrollment {
            remote_control_target: normalize_remote_control_url("http://localhost/backend-api/")
                .expect("target should normalize"),
            account_id: "account_id".to_string(),
            environment_id: "env_test".to_string(),
            server_id: "srv_e_test".to_string(),
            server_name: "test-server".to_string(),
            remote_control_token: remote_control_token.map(str::to_string),
            expires_at: remote_control_token
                .map(|_| time::OffsetDateTime::now_utc() + time::Duration::hours(1)),
        }
    }

    fn test_current_enrollment(
        enrollment: Option<RemoteControlEnrollment>,
    ) -> CurrentRemoteControlEnrollment {
        Arc::new(RemoteControlEnrollmentState::new(enrollment))
    }

    #[test]
    fn next_reconnect_delay_resets_after_cap() {
        let mut reconnect_attempt = 9;

        let (reconnect_delay, reconnect_backoff_reset) =
            next_reconnect_delay(&mut reconnect_attempt);

        assert_eq!(reconnect_delay, REMOTE_CONTROL_RECONNECT_BACKOFF_CAP);
        assert!(reconnect_backoff_reset);
        assert_eq!(reconnect_attempt, 0);

        let (reconnect_delay, reconnect_backoff_reset) =
            next_reconnect_delay(&mut reconnect_attempt);

        assert!(reconnect_delay >= Duration::from_millis(180));
        assert!(reconnect_delay <= Duration::from_millis(220));
        assert!(!reconnect_backoff_reset);
        assert_eq!(reconnect_attempt, 1);
    }

    #[test]
    fn websocket_404_only_reports_explicit_missing_remote_app_server() {
        let cases = [
            (
                Some(br#"{"detail":"Remote app server not found"}"#.to_vec()),
                true,
            ),
            (
                Some(br#" { "detail": "Remote app server not found", "extra": true } "#.to_vec()),
                true,
            ),
            (Some(br#"{"detail":"Not Found"}"#.to_vec()), false),
            (Some(b"Not Found".to_vec()), false),
            (Some(b"{".to_vec()), false),
            (Some(Vec::new()), false),
            (None, false),
        ];

        for (body, expected) in cases {
            let response = tungstenite::http::Response::builder()
                .status(/*status*/ 404)
                .body(body)
                .expect("response should build");
            assert_eq!(
                websocket_response_reports_missing_remote_app_server(&response),
                expected
            );
        }

        let response = tungstenite::http::Response::builder()
            .status(/*status*/ 503)
            .body(Some(
                br#"{"detail":"Remote app server not found"}"#.to_vec(),
            ))
            .expect("response should build");
        assert!(!websocket_response_reports_missing_remote_app_server(
            &response
        ));
    }

    fn remote_control_status_channel() -> (
        RemoteControlStatusPublisher,
        watch::Receiver<RemoteControlStatusChangedNotification>,
    ) {
        let (status_tx, status_rx) = watch::channel(RemoteControlStatusChangedNotification {
            status: RemoteControlConnectionStatus::Connecting,
            server_name: "test-server".to_string(),
            installation_id: TEST_INSTALLATION_ID.to_string(),
            environment_id: None,
        });
        (RemoteControlStatusPublisher::new(status_tx), status_rx)
    }

    fn enabled_desired_state_sender() -> watch::Sender<RemoteControlDesiredState> {
        watch::channel(RemoteControlDesiredState::Enabled {
            persistence_preference: None,
        })
        .0
    }

    #[test]
    fn mark_recovery_auth_change_seen_marks_only_recovery_revision_seen() {
        let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64);
        let auth_change_revision_before_recovery = *auth_change_rx.borrow();
        auth_change_tx.send_modify(|revision| *revision += 1);

        mark_recovery_auth_change_seen(&mut auth_change_rx, auth_change_revision_before_recovery);

        assert!(
            !auth_change_rx
                .has_changed()
                .expect("auth change watch should remain open")
        );
    }

    #[test]
    fn mark_recovery_auth_change_seen_preserves_racing_auth_change() {
        let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64);
        let auth_change_revision_before_recovery = *auth_change_rx.borrow();
        auth_change_tx.send_modify(|revision| *revision += 1);
        auth_change_tx.send_modify(|revision| *revision += 1);

        mark_recovery_auth_change_seen(&mut auth_change_rx, auth_change_revision_before_recovery);

        assert!(
            auth_change_rx
                .has_changed()
                .expect("auth change watch should remain open")
        );
    }

    async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc<StateRuntime> {
        StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string())
            .await
            .expect("state runtime should initialize")
    }

    fn remote_control_auth_manager() -> Arc<AuthManager> {
        auth_manager_from_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing())
    }

    fn remote_control_url_for_listener(listener: &TcpListener) -> String {
        let addr = listener
            .local_addr()
            .expect("listener should have a local addr");
        format!("http://{addr}/backend-api/")
    }

    fn remote_control_auth_dot_json(access_token: &str) -> AuthDotJson {
        #[derive(serde::Serialize)]
        struct Header {
            alg: &'static str,
            typ: &'static str,
        }

        let header = Header {
            alg: "none",
            typ: "JWT",
        };
        let payload = serde_json::json!({
            "email": "user@example.com",
            "https://api.openai.com/auth": {
                "chatgpt_user_id": "user-12345",
                "user_id": "user-12345",
                "chatgpt_account_id": "account_id"
            }
        });
        let b64 = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes);
        let header_b64 = b64(&serde_json::to_vec(&header).expect("header should serialize"));
        let payload_b64 = b64(&serde_json::to_vec(&payload).expect("payload should serialize"));
        let fake_jwt = format!("{header_b64}.{payload_b64}.sig");

        AuthDotJson {
            auth_mode: Some(AuthMode::Chatgpt),
            openai_api_key: None,
            tokens: Some(TokenData {
                id_token: parse_chatgpt_jwt_claims(&fake_jwt).expect("fake jwt should parse"),
                access_token: access_token.to_string(),
                refresh_token: "refresh-token".to_string(),
                account_id: Some("account_id".to_string()),
            }),
            last_refresh: Some(Utc::now()),
            agent_identity: None,
            personal_access_token: None,
            bedrock_api_key: None,
        }
    }

    #[tokio::test]
    async fn connect_remote_control_websocket_includes_http_error_details() {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("listener should bind");
        let remote_control_url = remote_control_url_for_listener(&listener);
        let remote_control_target =
            normalize_remote_control_url(&remote_control_url).expect("target should parse");
        let expected_error = format!(
            "failed to connect app-server remote control websocket `{}`: HTTP error: 503 Service Unavailable, request-id: <none>, cf-ray: <none>, body: upstream unavailable",
            remote_control_target.websocket_url
        );
        let server_task = tokio::spawn(async move {
            let (stream, request_line) = accept_http_request(&listener).await;
            assert_eq!(
                request_line,
                "GET /backend-api/wham/remote/control/server HTTP/1.1"
            );
            respond_with_status_and_headers(
                stream,
                "503 Service Unavailable",
                &[("x-trace-id", "trace-503"), ("x-region", "us-east-1")],
                "upstream unavailable",
            )
            .await;
        });
        let codex_home = TempDir::new().expect("temp dir should create");
        let state_db = remote_control_state_runtime(&codex_home).await;
        let auth_manager = remote_control_auth_manager();
        let mut auth_recovery = auth_manager.unauthorized_recovery();
        let mut auth_change_rx = auth_manager.auth_change_receiver();
        let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some(
            TEST_REMOTE_CONTROL_SERVER_TOKEN,
        ))));
        let (status_publisher, status_rx) = remote_control_status_channel();

        let err = match connect_remote_control_websocket(
            &remote_control_target,
            Some(state_db.as_ref()),
            RemoteControlAuthContext {
                auth_manager: &auth_manager,
                auth_recovery: &mut auth_recovery,
                auth_change_rx: &mut auth_change_rx,
            },
            &current_enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                server_name: "test-server",
                subscribe_cursor: None,
                app_server_client_name: None,
                desired_state_tx: &enabled_desired_state_sender(),
                desired_state_persistence_lock: &Semaphore::new(1),
            },
            &status_publisher,
        )
        .await
        {
            Ok(_) => panic!("http error response should fail the websocket connect"),
            Err(err) => err,
        };

        server_task.await.expect("server task should succeed");
        assert_eq!(err.to_string(), expected_error);
        assert!(current_enrollment.lock().await.is_some());
        assert_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Connecting,
                server_name: "test-server".to_string(),
                installation_id: TEST_INSTALLATION_ID.to_string(),
                environment_id: Some("env_test".to_string()),
            }
        );
    }

    #[tokio::test]
    async fn connect_remote_control_websocket_invalidates_unauthorized_server_token() {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("listener should bind");
        let remote_control_url = remote_control_url_for_listener(&listener);
        let remote_control_target =
            normalize_remote_control_url(&remote_control_url).expect("target should parse");
        let codex_home = TempDir::new().expect("temp dir should create");
        let state_db = remote_control_state_runtime(&codex_home).await;
        let auth_manager = remote_control_auth_manager();
        let mut auth_recovery = auth_manager.unauthorized_recovery();
        let mut auth_change_rx = auth_manager.auth_change_receiver();
        let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some(
            TEST_REMOTE_CONTROL_SERVER_TOKEN,
        ))));
        let (status_publisher, status_rx) = remote_control_status_channel();

        let server_task = tokio::spawn(async move {
            let (stream, request_line) = accept_http_request(&listener).await;
            assert_eq!(
                request_line,
                "GET /backend-api/wham/remote/control/server HTTP/1.1"
            );
            respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await;
        });

        let err = connect_remote_control_websocket(
            &remote_control_target,
            Some(state_db.as_ref()),
            RemoteControlAuthContext {
                auth_manager: &auth_manager,
                auth_recovery: &mut auth_recovery,
                auth_change_rx: &mut auth_change_rx,
            },
            &current_enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                server_name: "test-server",
                subscribe_cursor: None,
                app_server_client_name: None,
                desired_state_tx: &enabled_desired_state_sender(),
                desired_state_persistence_lock: &Semaphore::new(1),
            },
            &status_publisher,
        )
        .await
        .expect_err("unauthorized response should fail the websocket connect");

        server_task.await.expect("server task should succeed");
        assert_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Connecting,
                server_name: "test-server".to_string(),
                installation_id: TEST_INSTALLATION_ID.to_string(),
                environment_id: Some("env_test".to_string()),
            }
        );
        assert_eq!(
            err.to_string(),
            "remote control websocket auth failed with HTTP 401 Unauthorized; refreshing server token before reconnect"
        );
        let mut expected_enrollment = remote_control_enrollment(/*remote_control_token*/ None);
        expected_enrollment.remote_control_target = remote_control_target;
        assert_eq!(*current_enrollment.lock().await, Some(expected_enrollment));
    }

    #[tokio::test]
    async fn connect_remote_control_websocket_recovers_after_unauthorized_enrollment() {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("listener should bind");
        let remote_control_url = remote_control_url_for_listener(&listener);
        let remote_control_target =
            normalize_remote_control_url(&remote_control_url).expect("target should parse");
        let enroll_url = remote_control_target.enroll_url.clone();
        let server_task = tokio::spawn(async move {
            let (stream, request_line) = accept_http_request(&listener).await;
            assert_eq!(
                request_line,
                "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1"
            );
            respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await;
        });
        let codex_home = TempDir::new().expect("temp dir should create");
        save_auth(
            codex_home.path(),
            &remote_control_auth_dot_json("stale-token"),
            AuthCredentialsStoreMode::File,
            AuthKeyringBackendKind::default(),
        )
        .expect("stale auth should save");
        let state_db = remote_control_state_runtime(&codex_home).await;
        let auth_manager = AuthManager::shared(
            codex_home.path().to_path_buf(),
            /*enable_codex_api_key_env*/ false,
            AuthCredentialsStoreMode::File,
            /*chatgpt_base_url*/ None,
            AuthKeyringBackendKind::default(),
        )
        .await;
        let mut auth_recovery = auth_manager.unauthorized_recovery();
        let mut auth_change_rx = auth_manager.auth_change_receiver();
        let current_enrollment = test_current_enrollment(/*enrollment*/ None);
        let (status_publisher, status_rx) = remote_control_status_channel();
        save_auth(
            codex_home.path(),
            &remote_control_auth_dot_json("fresh-token"),
            AuthCredentialsStoreMode::File,
            AuthKeyringBackendKind::default(),
        )
        .expect("fresh auth should save");

        let err = connect_remote_control_websocket(
            &remote_control_target,
            Some(state_db.as_ref()),
            RemoteControlAuthContext {
                auth_manager: &auth_manager,
                auth_recovery: &mut auth_recovery,
                auth_change_rx: &mut auth_change_rx,
            },
            &current_enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                server_name: "test-server",
                subscribe_cursor: None,
                app_server_client_name: None,
                desired_state_tx: &enabled_desired_state_sender(),
                desired_state_persistence_lock: &Semaphore::new(1),
            },
            &status_publisher,
        )
        .await
        .expect_err("unauthorized enrollment should fail the websocket connect");

        server_task.await.expect("server task should succeed");
        assert!(
            !status_rx
                .has_changed()
                .expect("remote control status watch should remain open")
        );
        assert_eq!(
            err.to_string(),
            format!(
                "remote control server enrollment failed at `{enroll_url}`: HTTP 401 Unauthorized, request-id: <none>, cf-ray: <none>, body: unauthorized; retrying after auth recovery"
            )
        );
        assert_eq!(
            auth_manager
                .auth()
                .await
                .expect("auth should remain available")
                .get_token()
                .expect("token should be readable"),
            "fresh-token"
        );
        assert!(
            !auth_change_rx
                .has_changed()
                .expect("auth change watch should remain open"),
            "recovery's own auth reload should not wake the reconnect loop"
        );
    }

    #[tokio::test]
    async fn connect_remote_control_websocket_recovers_after_unauthorized_refresh() {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("listener should bind");
        let remote_control_url = remote_control_url_for_listener(&listener);
        let remote_control_target =
            normalize_remote_control_url(&remote_control_url).expect("target should parse");
        let refresh_url = remote_control_target.refresh_url.clone();
        let server_task = tokio::spawn(async move {
            let (stream, request_line) = accept_http_request(&listener).await;
            assert_eq!(
                request_line,
                "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1"
            );
            respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await;
        });
        let codex_home = TempDir::new().expect("temp dir should create");
        save_auth(
            codex_home.path(),
            &remote_control_auth_dot_json("stale-token"),
            AuthCredentialsStoreMode::File,
            AuthKeyringBackendKind::default(),
        )
        .expect("stale auth should save");
        let state_db = remote_control_state_runtime(&codex_home).await;
        let auth_manager = AuthManager::shared(
            codex_home.path().to_path_buf(),
            /*enable_codex_api_key_env*/ false,
            AuthCredentialsStoreMode::File,
            /*chatgpt_base_url*/ None,
            AuthKeyringBackendKind::default(),
        )
        .await;
        let mut auth_recovery = auth_manager.unauthorized_recovery();
        let mut auth_change_rx = auth_manager.auth_change_receiver();
        let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(
            /*remote_control_token*/ None,
        )));
        let (status_publisher, status_rx) = remote_control_status_channel();
        save_auth(
            codex_home.path(),
            &remote_control_auth_dot_json("fresh-token"),
            AuthCredentialsStoreMode::File,
            AuthKeyringBackendKind::default(),
        )
        .expect("fresh auth should save");

        let err = connect_remote_control_websocket(
            &remote_control_target,
            Some(state_db.as_ref()),
            RemoteControlAuthContext {
                auth_manager: &auth_manager,
                auth_recovery: &mut auth_recovery,
                auth_change_rx: &mut auth_change_rx,
            },
            &current_enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                server_name: "test-server",
                subscribe_cursor: None,
                app_server_client_name: None,
                desired_state_tx: &enabled_desired_state_sender(),
                desired_state_persistence_lock: &Semaphore::new(1),
            },
            &status_publisher,
        )
        .await
        .expect_err("unauthorized refresh should fail the websocket connect");

        server_task.await.expect("server task should succeed");
        assert_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Connecting,
                server_name: "test-server".to_string(),
                installation_id: TEST_INSTALLATION_ID.to_string(),
                environment_id: Some("env_test".to_string()),
            }
        );
        assert_eq!(
            err.to_string(),
            format!(
                "remote control server refresh failed at `{refresh_url}`: HTTP 401 Unauthorized, request-id: <none>, cf-ray: <none>, body: unauthorized; retrying after auth recovery"
            )
        );
        assert_eq!(
            auth_manager
                .auth()
                .await
                .expect("auth should remain available")
                .get_token()
                .expect("token should be readable"),
            "fresh-token"
        );
        assert!(
            !auth_change_rx
                .has_changed()
                .expect("auth change watch should remain open"),
            "recovery's own auth reload should not wake the reconnect loop"
        );
    }

    #[tokio::test]
    async fn connect_remote_control_websocket_requires_sqlite_state_db() {
        let remote_control_target = normalize_remote_control_url("http://127.0.0.1:9/backend-api/")
            .expect("target should parse");
        let auth_manager = remote_control_auth_manager();
        let mut auth_recovery = auth_manager.unauthorized_recovery();
        let mut auth_change_rx = auth_manager.auth_change_receiver();
        let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some(
            TEST_REMOTE_CONTROL_SERVER_TOKEN,
        ))));
        let (status_publisher, _status_rx) = remote_control_status_channel();

        let err = connect_remote_control_websocket(
            &remote_control_target,
            /*state_db*/ None,
            RemoteControlAuthContext {
                auth_manager: &auth_manager,
                auth_recovery: &mut auth_recovery,
                auth_change_rx: &mut auth_change_rx,
            },
            &current_enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                server_name: "test-server",
                subscribe_cursor: None,
                app_server_client_name: None,
                desired_state_tx: &enabled_desired_state_sender(),
                desired_state_persistence_lock: &Semaphore::new(1),
            },
            &status_publisher,
        )
        .await
        .expect_err("missing sqlite state db should fail remote control");

        assert_eq!(err.kind(), ErrorKind::NotFound);
        assert_eq!(err.to_string(), "remote control requires sqlite state db");
        assert_eq!(*current_enrollment.lock().await, None);
    }

    #[tokio::test]
    async fn connect_remote_control_websocket_requires_chatgpt_auth() {
        let remote_control_target = normalize_remote_control_url("http://127.0.0.1:9/backend-api/")
            .expect("target should parse");
        let codex_home = TempDir::new().expect("temp dir should create");
        let state_db = remote_control_state_runtime(&codex_home).await;
        let auth_manager = AuthManager::shared(
            codex_home.path().to_path_buf(),
            /*enable_codex_api_key_env*/ false,
            AuthCredentialsStoreMode::File,
            /*chatgpt_base_url*/ None,
            AuthKeyringBackendKind::default(),
        )
        .await;
        let mut auth_recovery = auth_manager.unauthorized_recovery();
        let mut auth_change_rx = auth_manager.auth_change_receiver();
        let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some(
            TEST_REMOTE_CONTROL_SERVER_TOKEN,
        ))));
        let (status_publisher, mut status_rx) = remote_control_status_channel();
        status_publisher.publish_environment_id(Some("env_test".to_string()));
        status_rx
            .changed()
            .await
            .expect("remote control status watch should remain open");

        let err = connect_remote_control_websocket(
            &remote_control_target,
            Some(state_db.as_ref()),
            RemoteControlAuthContext {
                auth_manager: &auth_manager,
                auth_recovery: &mut auth_recovery,
                auth_change_rx: &mut auth_change_rx,
            },
            &current_enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                server_name: "test-server",
                subscribe_cursor: None,
                app_server_client_name: None,
                desired_state_tx: &enabled_desired_state_sender(),
                desired_state_persistence_lock: &Semaphore::new(1),
            },
            &status_publisher,
        )
        .await
        .expect_err("missing auth should fail remote control");

        assert_eq!(err.kind(), ErrorKind::PermissionDenied);
        assert_eq!(
            err.to_string(),
            "remote control requires ChatGPT authentication"
        );
        assert_eq!(*current_enrollment.lock().await, None);
        assert_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Connecting,
                server_name: "test-server".to_string(),
                installation_id: TEST_INSTALLATION_ID.to_string(),
                environment_id: None,
            }
        );
    }

    #[tokio::test]
    async fn run_remote_control_websocket_loop_shutdown_cancels_reconnect_backoff() {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("listener should bind");
        let remote_control_url = remote_control_url_for_listener(&listener);
        drop(listener);

        let remote_control_target =
            normalize_remote_control_url(&remote_control_url).expect("target should parse");
        let (transport_event_tx, transport_event_rx) = mpsc::channel(1);
        drop(transport_event_rx);
        let (status_publisher, _status_rx) = remote_control_status_channel();
        let shutdown_token = CancellationToken::new();
        let (desired_state_tx, _desired_state_rx) =
            watch::channel(RemoteControlDesiredState::Enabled {
                persistence_preference: None,
            });
        let websocket_task = tokio::spawn({
            let shutdown_token = shutdown_token.clone();
            async move {
                RemoteControlWebsocket::new(
                    RemoteControlWebsocketConfig {
                        remote_control_url,
                        installation_id: TEST_INSTALLATION_ID.to_string(),
                        remote_control_target: Some(remote_control_target),
                        server_name: "test-server".to_string(),
                    },
                    /*state_db*/ None,
                    remote_control_auth_manager(),
                    RemoteControlChannels {
                        transport_event_tx,
                        status_publisher,
                        current_enrollment: test_current_enrollment(/*enrollment*/ None),
                        pairing_persistence_key: watch::channel(None).0,
                        desired_state_persistence_lock: Arc::new(Semaphore::new(1)),
                    },
                    shutdown_token,
                    Arc::new(desired_state_tx),
                )
                .run(/*app_server_client_name_rx*/ None)
                .await
            }
        });

        tokio::time::sleep(Duration::from_millis(50)).await;
        shutdown_token.cancel();

        timeout(Duration::from_millis(100), websocket_task)
            .await
            .expect("shutdown should cancel reconnect backoff")
            .expect("websocket task should join");
    }

    #[tokio::test]
    async fn publish_status_if_changed_sends_only_status_changes() {
        let (status_publisher, mut status_rx) = remote_control_status_channel();

        status_publisher.publish_environment_id(/*environment_id*/ None);
        assert!(
            timeout(Duration::from_millis(20), status_rx.changed())
                .await
                .is_err()
        );

        status_publisher.publish_environment_id(Some("env_first".to_string()));
        status_rx
            .changed()
            .await
            .expect("remote control status watch should remain open");
        assert_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Connecting,
                server_name: "test-server".to_string(),
                installation_id: TEST_INSTALLATION_ID.to_string(),
                environment_id: Some("env_first".to_string()),
            }
        );

        status_publisher.publish_environment_id(Some("env_first".to_string()));
        assert!(
            timeout(Duration::from_millis(20), status_rx.changed())
                .await
                .is_err()
        );

        status_publisher.publish_status(RemoteControlConnectionStatus::Connected);
        status_rx
            .changed()
            .await
            .expect("remote control status watch should remain open");
        assert_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Connected,
                server_name: "test-server".to_string(),
                installation_id: TEST_INSTALLATION_ID.to_string(),
                environment_id: Some("env_first".to_string()),
            }
        );

        status_publisher.publish_environment_id(/*environment_id*/ None);
        status_rx
            .changed()
            .await
            .expect("remote control status watch should remain open");
        assert_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Connected,
                server_name: "test-server".to_string(),
                installation_id: TEST_INSTALLATION_ID.to_string(),
                environment_id: None,
            }
        );

        status_publisher.publish_environment_id(Some("env_disabled".to_string()));
        status_publisher.publish_status(RemoteControlConnectionStatus::Disabled);
        status_rx
            .changed()
            .await
            .expect("remote control status watch should remain open");
        assert_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Disabled,
                server_name: "test-server".to_string(),
                installation_id: TEST_INSTALLATION_ID.to_string(),
                environment_id: None,
            }
        );

        status_publisher.publish_environment_id(Some("env_disabled".to_string()));
        assert!(
            timeout(Duration::from_millis(20), status_rx.changed())
                .await
                .is_err()
        );
    }

    #[tokio::test]
    async fn run_server_writer_inner_sends_periodic_ping_frames() {
        let (client_stream, mut server_stream) = connected_websocket_pair().await;
        let (websocket_writer, _websocket_reader) = client_stream.split();
        let (outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
        let state = Arc::new(Mutex::new(WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        }));
        let (_server_event_tx, server_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY);
        let server_event_rx = Arc::new(Mutex::new(server_event_rx));
        let shutdown_token = CancellationToken::new();
        let writer_task = tokio::spawn(RemoteControlWebsocket::run_server_writer_inner(
            state,
            server_event_rx,
            used_rx,
            websocket_writer,
            Duration::from_millis(20),
            shutdown_token.clone(),
        ));

        let message = timeout(Duration::from_secs(5), server_stream.next())
            .await
            .expect("ping frame should arrive in time")
            .expect("server websocket should stay open")
            .expect("ping frame should read");
        assert!(matches!(message, tungstenite::Message::Ping(_)));

        shutdown_token.cancel();
        writer_task
            .await
            .expect("writer task should join")
            .expect("writer should stop cleanly");
    }

    #[tokio::test]
    async fn join_connection_workers_aborts_stuck_worker_after_timeout() {
        let mut join_set = tokio::task::JoinSet::new();
        join_set.spawn(futures::future::pending::<()>());

        RemoteControlWebsocket::join_connection_workers(&mut join_set, Duration::from_millis(10))
            .await;

        assert!(join_set.is_empty());
    }

    #[tokio::test]
    async fn run_server_writer_inner_assigns_contiguous_seq_ids_per_stream() {
        let (client_stream, mut server_stream) = connected_websocket_pair().await;
        let (websocket_writer, _websocket_reader) = client_stream.split();
        let (outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
        let state = Arc::new(Mutex::new(WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        }));
        let (server_event_tx, server_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY);
        let server_event_rx = Arc::new(Mutex::new(server_event_rx));
        let shutdown_token = CancellationToken::new();
        let writer_task = tokio::spawn(RemoteControlWebsocket::run_server_writer_inner(
            state,
            server_event_rx,
            used_rx,
            websocket_writer,
            Duration::from_secs(60),
            shutdown_token.clone(),
        ));

        let client_id = ClientId("client-1".to_string());
        let first_stream = StreamId("stream-1".to_string());
        let second_stream = StreamId("stream-2".to_string());
        for stream_id in [&first_stream, &second_stream, &first_stream] {
            server_event_tx
                .send(super::super::QueuedServerEnvelope {
                    event: ServerEvent::Pong {
                        status: crate::transport::remote_control::protocol::PongStatus::Active,
                    },
                    client_id: client_id.clone(),
                    stream_id: stream_id.clone(),
                    write_complete_tx: None,
                })
                .await
                .expect("server event should queue");
        }

        assert_eq!(
            read_server_text_event(&mut server_stream).await,
            serde_json::json!({
                "type": "pong",
                "client_id": "client-1",
                "stream_id": "stream-1",
                "seq_id": 1,
                "status": "active",
            })
        );
        assert_eq!(
            read_server_text_event(&mut server_stream).await,
            serde_json::json!({
                "type": "pong",
                "client_id": "client-1",
                "stream_id": "stream-2",
                "seq_id": 1,
                "status": "active",
            })
        );
        assert_eq!(
            read_server_text_event(&mut server_stream).await,
            serde_json::json!({
                "type": "pong",
                "client_id": "client-1",
                "stream_id": "stream-1",
                "seq_id": 2,
                "status": "active",
            })
        );

        shutdown_token.cancel();
        writer_task
            .await
            .expect("writer task should join")
            .expect("writer should stop cleanly");
    }

    #[tokio::test]
    async fn run_websocket_reader_inner_times_out_without_pong_frames() {
        let (client_stream, _server_stream) = connected_websocket_pair().await;
        let (_websocket_writer, websocket_reader) = client_stream.split();
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let state = Arc::new(Mutex::new(WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        }));
        let (server_event_tx, _server_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY);
        let (transport_event_tx, _transport_event_rx) =
            mpsc::channel(super::super::CHANNEL_CAPACITY);
        let shutdown_token = CancellationToken::new();
        let client_tracker = Arc::new(Mutex::new(ClientTracker::new(
            server_event_tx,
            transport_event_tx,
            &shutdown_token,
        )));

        let err = timeout(
            Duration::from_secs(5),
            RemoteControlWebsocket::run_websocket_reader_inner(
                client_tracker,
                state,
                websocket_reader,
                Duration::from_millis(100),
                shutdown_token,
            ),
        )
        .await
        .expect("reader should time out waiting for pong")
        .expect_err("missing pong should fail the websocket reader");

        assert_eq!(err.kind(), ErrorKind::TimedOut);
        assert_eq!(err.to_string(), "remote control websocket pong timeout");
    }

    #[test]
    fn outbound_buffer_acks_by_stream_id() {
        let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
        let client_1 = ClientId("client-1".to_string());
        let client_2 = ClientId("client-2".to_string());
        let stream_1 = StreamId("stream-1".to_string());

        outbound_buffer.insert(&server_envelope(
            &client_1,
            "stream-1",
            /*seq_id*/ 1,
            "first-client-old-stream",
        ));
        outbound_buffer.insert(&server_envelope(
            &client_2,
            "stream-1",
            /*seq_id*/ 2,
            "second-client",
        ));
        outbound_buffer.insert(&server_envelope(
            &client_1,
            "stream-2",
            /*seq_id*/ 3,
            "first-client-new-stream",
        ));

        outbound_buffer.ack(
            &client_1, &stream_1, /*acked_seq_id*/ 3, /*acked_segment_id*/ None,
        );

        let mut retained = outbound_buffer
            .server_envelopes()
            .map(|server_envelope| {
                (
                    server_envelope.client_id.0.as_str(),
                    server_envelope.stream_id.0.as_str(),
                    server_envelope.seq_id,
                )
            })
            .collect::<Vec<_>>();
        retained.sort_unstable();
        assert_eq!(
            retained,
            vec![("client-1", "stream-2", 3), ("client-2", "stream-1", 2)]
        );
        assert_eq!(*used_rx.borrow(), 2);
    }

    #[test]
    fn outbound_buffer_retains_unacked_messages_until_ack_advances() {
        let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
        let client_1 = ClientId("client-1".to_string());
        let client_2 = ClientId("client-2".to_string());
        let stream_1 = StreamId("stream-1".to_string());

        outbound_buffer.insert(&server_envelope(
            &client_1,
            "stream-1",
            /*seq_id*/ 1,
            "first-old",
        ));
        outbound_buffer.insert(&server_envelope(
            &client_1,
            "stream-2",
            /*seq_id*/ 2,
            "first-new",
        ));
        outbound_buffer.insert(&server_envelope(
            &client_2, "stream-1", /*seq_id*/ 3, "second",
        ));

        outbound_buffer.ack(
            &client_1, &stream_1, /*acked_seq_id*/ 1, /*acked_segment_id*/ None,
        );

        let mut retained = outbound_buffer
            .server_envelopes()
            .map(|server_envelope| {
                (
                    server_envelope.client_id.0.as_str(),
                    server_envelope.stream_id.0.as_str(),
                    server_envelope.seq_id,
                )
            })
            .collect::<Vec<_>>();
        retained.sort_unstable();
        assert_eq!(
            retained,
            vec![("client-1", "stream-2", 2), ("client-2", "stream-1", 3)]
        );
        assert_eq!(*used_rx.borrow(), 2);
    }

    #[test]
    fn outbound_buffer_advances_segmented_acks_by_wire_cursor() {
        let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
        let client_id = ClientId("client-1".to_string());
        let stream_id = StreamId("stream-1".to_string());

        outbound_buffer.insert(&server_chunk_envelope(
            &client_id, "stream-1", /*seq_id*/ 4, /*segment_id*/ 0,
        ));
        outbound_buffer.insert(&server_chunk_envelope(
            &client_id, "stream-1", /*seq_id*/ 4, /*segment_id*/ 1,
        ));

        outbound_buffer.ack(
            &client_id,
            &stream_id,
            /*acked_seq_id*/ 4,
            /*acked_segment_id*/ Some(1),
        );

        let retained = outbound_buffer
            .server_envelopes()
            .map(|server_envelope| server_envelope.event.segment_id())
            .collect::<Vec<_>>();
        assert_eq!(retained, Vec::<Option<usize>>::new());
        assert_eq!(*used_rx.borrow(), 0);
    }

    #[test]
    fn outbound_buffer_treats_segmentless_acks_as_seq_level_acks() {
        let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new();
        let client_id = ClientId("client-1".to_string());
        let stream_id = StreamId("stream-1".to_string());

        outbound_buffer.insert(&server_chunk_envelope(
            &client_id, "stream-1", /*seq_id*/ 4, /*segment_id*/ 0,
        ));
        outbound_buffer.insert(&server_chunk_envelope(
            &client_id, "stream-1", /*seq_id*/ 4, /*segment_id*/ 1,
        ));

        outbound_buffer.ack(
            &client_id, &stream_id, /*acked_seq_id*/ 4, /*acked_segment_id*/ None,
        );

        let retained = outbound_buffer
            .server_envelopes()
            .map(|server_envelope| server_envelope.event.segment_id())
            .collect::<Vec<_>>();
        assert_eq!(retained, Vec::<Option<usize>>::new());
        assert_eq!(*used_rx.borrow(), 0);
    }

    #[test]
    fn websocket_state_drops_duplicate_client_chunks_while_pending() {
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let mut state = WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        };
        let first_chunk = client_chunk_envelope(
            "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0,
            /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x",
        );
        let second_chunk = client_chunk_envelope(
            "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 1,
            /*segment_count*/ 2, /*message_size_bytes*/ 2, b"y",
        );

        assert!(matches!(
            observe_client_message(&mut state, first_chunk.clone()),
            ClientSegmentObservation::Pending
        ));
        assert!(matches!(
            observe_client_message(&mut state, first_chunk.clone()),
            ClientSegmentObservation::Dropped
        ));
        assert!(matches!(
            observe_client_message(&mut state, second_chunk),
            ClientSegmentObservation::Dropped
        ));
        assert!(matches!(
            observe_client_message(&mut state, first_chunk),
            ClientSegmentObservation::Pending
        ));
    }

    #[test]
    fn websocket_state_drops_replayed_client_chunks_after_completion() {
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let mut state = WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        };
        let message = JSONRPCMessage::Notification(JSONRPCNotification {
            method: "initialized".to_string(),
            params: None,
        });
        let raw = serde_json::to_vec(&message).expect("message should serialize");
        let split = raw.len() / 2;
        let first_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 4,
            /*segment_id*/ 0,
            /*segment_count*/ 2,
            raw.len(),
            &raw[..split],
        );
        let second_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 4,
            /*segment_id*/ 1,
            /*segment_count*/ 2,
            raw.len(),
            &raw[split..],
        );

        assert!(matches!(
            observe_client_message(&mut state, first_chunk.clone()),
            ClientSegmentObservation::Pending
        ));
        let completed_envelope = match observe_client_message(&mut state, second_chunk) {
            ClientSegmentObservation::Forward(client_envelope) => *client_envelope,
            _ => panic!("expected completed client message"),
        };
        state.record_client_message_delivery(
            &completed_envelope,
            Some((
                (
                    ClientId("client-1".to_string()),
                    Some(StreamId("stream-1".to_string())),
                ),
                4,
            )),
        );
        assert!(matches!(
            observe_client_message(&mut state, first_chunk),
            ClientSegmentObservation::Dropped
        ));
    }

    #[test]
    fn websocket_state_allows_replay_before_completed_chunk_delivery() {
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let mut state = WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        };
        let message = JSONRPCMessage::Notification(JSONRPCNotification {
            method: "initialized".to_string(),
            params: None,
        });
        let raw = serde_json::to_vec(&message).expect("message should serialize");
        let split = raw.len() / 2;
        let first_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 4,
            /*segment_id*/ 0,
            /*segment_count*/ 2,
            raw.len(),
            &raw[..split],
        );
        let second_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 4,
            /*segment_id*/ 1,
            /*segment_count*/ 2,
            raw.len(),
            &raw[split..],
        );

        assert!(matches!(
            observe_client_message(&mut state, first_chunk.clone()),
            ClientSegmentObservation::Pending
        ));
        assert!(matches!(
            observe_client_message(&mut state, second_chunk),
            ClientSegmentObservation::Forward(_)
        ));
        assert!(matches!(
            observe_client_message(&mut state, first_chunk),
            ClientSegmentObservation::Pending
        ));
    }

    #[test]
    fn websocket_state_allows_replay_after_rejected_out_of_order_chunk() {
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let mut state = WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        };
        let first_chunk = client_chunk_envelope(
            "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0,
            /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x",
        );
        let second_chunk = client_chunk_envelope(
            "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 1,
            /*segment_count*/ 2, /*message_size_bytes*/ 2, b"y",
        );

        assert!(matches!(
            observe_client_message(&mut state, second_chunk),
            ClientSegmentObservation::Dropped
        ));
        assert!(matches!(
            observe_client_message(&mut state, first_chunk),
            ClientSegmentObservation::Pending
        ));
    }

    #[test]
    fn websocket_state_allows_replay_after_later_chunk_drops() {
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let mut state = WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        };
        let first_chunk = client_chunk_envelope(
            "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0,
            /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x",
        );
        let invalid_second_chunk = client_chunk_envelope(
            "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 1,
            /*segment_count*/ 2, /*message_size_bytes*/ 2, b"",
        );

        assert!(matches!(
            observe_client_message(&mut state, first_chunk.clone()),
            ClientSegmentObservation::Pending
        ));
        assert!(matches!(
            observe_client_message(&mut state, invalid_second_chunk),
            ClientSegmentObservation::Dropped
        ));
        assert!(matches!(
            observe_client_message(&mut state, first_chunk),
            ClientSegmentObservation::Pending
        ));
    }

    #[test]
    fn websocket_state_drops_oversized_client_chunk_frames() {
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let mut state = WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        };
        let chunk = client_chunk_envelope(
            "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0,
            /*segment_count*/ 1, /*message_size_bytes*/ 1, b"x",
        );

        assert!(matches!(
            state.observe_client_message(chunk, REMOTE_CONTROL_SEGMENT_MAX_BYTES + 1),
            ClientSegmentObservation::Dropped
        ));
    }

    #[test]
    fn websocket_state_ignores_oversized_stale_chunks_without_dropping_newer_assembly() {
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let mut state = WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        };
        let message = JSONRPCMessage::Notification(JSONRPCNotification {
            method: "initialized".to_string(),
            params: None,
        });
        let raw = serde_json::to_vec(&message).expect("message should serialize");
        let split = raw.len() / 2;
        let first_newer_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 8,
            /*segment_id*/ 0,
            /*segment_count*/ 2,
            raw.len(),
            &raw[..split],
        );
        let oversized_stale_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 7,
            /*segment_id*/ 0,
            /*segment_count*/ 2,
            raw.len(),
            &raw[..split],
        );
        let second_newer_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 8,
            /*segment_id*/ 1,
            /*segment_count*/ 2,
            raw.len(),
            &raw[split..],
        );

        assert!(matches!(
            observe_client_message(&mut state, first_newer_chunk),
            ClientSegmentObservation::Pending
        ));
        assert!(matches!(
            state.observe_client_message(
                oversized_stale_chunk,
                REMOTE_CONTROL_SEGMENT_MAX_BYTES + 1,
            ),
            ClientSegmentObservation::Dropped
        ));
        assert!(matches!(
            observe_client_message(&mut state, second_newer_chunk),
            ClientSegmentObservation::Forward(_)
        ));
    }

    #[test]
    fn websocket_state_ignores_oversized_duplicate_chunks_without_dropping_current_assembly() {
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let mut state = WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        };
        let message = JSONRPCMessage::Notification(JSONRPCNotification {
            method: "initialized".to_string(),
            params: None,
        });
        let raw = serde_json::to_vec(&message).expect("message should serialize");
        let split = raw.len() / 2;
        let first_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 8,
            /*segment_id*/ 0,
            /*segment_count*/ 2,
            raw.len(),
            &raw[..split],
        );
        let oversized_duplicate_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 8,
            /*segment_id*/ 0,
            /*segment_count*/ 2,
            raw.len(),
            &raw[..split],
        );
        let second_chunk = client_chunk_envelope(
            "client-1",
            "stream-1",
            /*seq_id*/ 8,
            /*segment_id*/ 1,
            /*segment_count*/ 2,
            raw.len(),
            &raw[split..],
        );

        assert!(matches!(
            observe_client_message(&mut state, first_chunk),
            ClientSegmentObservation::Pending
        ));
        assert!(matches!(
            state.observe_client_message(
                oversized_duplicate_chunk,
                REMOTE_CONTROL_SEGMENT_MAX_BYTES + 1,
            ),
            ClientSegmentObservation::Dropped
        ));
        assert!(matches!(
            observe_client_message(&mut state, second_chunk),
            ClientSegmentObservation::Forward(_)
        ));
    }

    #[test]
    fn websocket_state_clears_chunk_cursor_when_stream_is_invalidated() {
        let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new();
        let mut state = WebsocketState {
            outbound_buffer,
            subscribe_cursor: None,
            next_seq_id_by_stream: HashMap::new(),
            last_completed_client_chunk_seq_id_by_stream: HashMap::new(),
            client_segment_reassembler: ClientSegmentReassembler::default(),
        };
        let client_id = ClientId("client-1".to_string());
        let stream_id = StreamId("stream-1".to_string());

        assert!(matches!(
            observe_client_message(
                &mut state,
                client_chunk_envelope(
                    "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0,
                    /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x",
                )
            ),
            ClientSegmentObservation::Pending
        ));
        state.invalidate_client_message_stream(&client_id, &stream_id);
        state
            .client_segment_reassembler
            .invalidate_stream(&client_id, &stream_id);

        assert!(matches!(
            observe_client_message(
                &mut state,
                client_chunk_envelope(
                    "client-1", "stream-1", /*seq_id*/ 1, /*segment_id*/ 0,
                    /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x",
                )
            ),
            ClientSegmentObservation::Pending
        ));
    }

    fn server_envelope(
        client_id: &ClientId,
        stream_id: &str,
        seq_id: u64,
        summary: &str,
    ) -> ServerEnvelope {
        ServerEnvelope {
            event: ServerEvent::ServerMessage {
                message: Box::new(OutgoingMessage::AppServerNotification(
                    ServerNotification::ConfigWarning(ConfigWarningNotification {
                        summary: summary.to_string(),
                        details: None,
                        path: None,
                        range: None,
                    }),
                )),
            },
            client_id: client_id.clone(),
            stream_id: StreamId(stream_id.to_string()),
            seq_id,
        }
    }

    fn server_chunk_envelope(
        client_id: &ClientId,
        stream_id: &str,
        seq_id: u64,
        segment_id: usize,
    ) -> ServerEnvelope {
        ServerEnvelope {
            event: ServerEvent::ServerMessageChunk {
                segment_id,
                segment_count: 2,
                message_size_bytes: 2,
                message_chunk_base64: String::new(),
            },
            client_id: client_id.clone(),
            stream_id: StreamId(stream_id.to_string()),
            seq_id,
        }
    }

    fn client_chunk_envelope(
        client_id: &str,
        stream_id: &str,
        seq_id: u64,
        segment_id: usize,
        segment_count: usize,
        message_size_bytes: usize,
        chunk: &[u8],
    ) -> ClientEnvelope {
        ClientEnvelope {
            event: ClientEvent::ClientMessageChunk {
                segment_id,
                segment_count,
                message_size_bytes,
                message_chunk_base64: base64::engine::general_purpose::STANDARD.encode(chunk),
            },
            client_id: ClientId(client_id.to_string()),
            stream_id: Some(StreamId(stream_id.to_string())),
            seq_id: Some(seq_id),
            cursor: None,
        }
    }

    fn observe_client_message(
        state: &mut WebsocketState,
        envelope: ClientEnvelope,
    ) -> ClientSegmentObservation {
        let wire_size_bytes = serde_json::to_vec(&envelope)
            .expect("client envelope should serialize")
            .len();
        state.observe_client_message(envelope, wire_size_bytes)
    }

    async fn accept_http_request(listener: &TcpListener) -> (TcpStream, String) {
        let (stream, _) = timeout(TEST_HTTP_ACCEPT_TIMEOUT, listener.accept())
            .await
            .expect("HTTP request should arrive in time")
            .expect("listener accept should succeed");
        let mut reader = BufReader::new(stream);

        let mut request_line = String::new();
        reader
            .read_line(&mut request_line)
            .await
            .expect("request line should read");
        loop {
            let mut line = String::new();
            reader
                .read_line(&mut line)
                .await
                .expect("header line should read");
            if line == "\r\n" {
                break;
            }
        }

        (
            reader.into_inner(),
            request_line.trim_end_matches("\r\n").to_string(),
        )
    }

    async fn connected_websocket_pair() -> (
        WebSocketStream<MaybeTlsStream<TcpStream>>,
        WebSocketStream<TcpStream>,
    ) {
        let listener = TcpListener::bind("127.0.0.1:0")
            .await
            .expect("listener should bind");
        let connect_task = tokio::spawn(connect_async(format!(
            "ws://{}",
            listener
                .local_addr()
                .expect("listener should have a local addr")
        )));
        let (server_stream, _) = listener
            .accept()
            .await
            .expect("server should accept client");
        let server_stream = accept_async(server_stream)
            .await
            .expect("server websocket handshake should succeed");
        let (client_stream, _) = connect_task
            .await
            .expect("client connect task should join")
            .expect("client websocket handshake should succeed");

        (client_stream, server_stream)
    }

    async fn read_server_text_event(
        server_stream: &mut WebSocketStream<TcpStream>,
    ) -> serde_json::Value {
        let message = timeout(Duration::from_secs(5), server_stream.next())
            .await
            .expect("server event should arrive in time")
            .expect("server websocket should stay open")
            .expect("server event should read");
        let tungstenite::Message::Text(text) = message else {
            panic!("expected text event, got {message:?}");
        };
        serde_json::from_str(text.as_ref()).expect("server event should deserialize")
    }

    async fn respond_with_status_and_headers(
        mut stream: TcpStream,
        status: &str,
        headers: &[(&str, &str)],
        body: &str,
    ) {
        let extra_headers = headers
            .iter()
            .map(|(name, value)| format!("{name}: {value}\r\n"))
            .collect::<String>();
        let response = format!(
            "HTTP/1.1 {status}\r\ncontent-type: text/plain\r\ncontent-length: {}\r\nconnection: close\r\n{extra_headers}\r\n{body}",
            body.len(),
        );
        stream
            .write_all(response.as_bytes())
            .await
            .expect("response should write");
        stream.flush().await.expect("response should flush");
    }
}
