use crate::transport::TransportEvent;
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::RemoteControlConnectionAuth;
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::update_persisted_remote_control_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::segment::ClientSegmentObservation;
use super::segment::ClientSegmentReassembler;
use super::segment::REMOTE_CONTROL_SEGMENT_MAX_BYTES;
use super::segment::split_server_envelope_for_transport;
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::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;
use tracing::error;
use tracing::info;
use tracing::warn;

pub(super) const REMOTE_CONTROL_PROTOCOL_VERSION: &str = "3";
pub(super) const REMOTE_CONTROL_ACCOUNT_ID_HEADER: &str = "chatgpt-account-id";
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);

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;
        }

        let observation = self.client_segment_reassembler.observe(client_envelope);
        if matches!(observation, ClientSegmentObservation::Forward(_))
            && let Some((key, seq_id)) = client_message_key
        {
            self.last_completed_client_chunk_seq_id_by_stream
                .insert(key, seq_id);
        }
        observation
    }

    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,
    remote_control_target: Option<RemoteControlTarget>,
    state_db: Option<Arc<StateRuntime>>,
    auth_manager: Arc<AuthManager>,
    status_publisher: RemoteControlStatusPublisher,
    shutdown_token: CancellationToken,
    reconnect_attempt: u64,
    enrollment: Option<RemoteControlEnrollment>,
    auth_recovery: UnauthorizedRecovery,
    client_tracker: Arc<Mutex<ClientTracker>>,
    state: Arc<Mutex<WebsocketState>>,
    server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>,
    used_rx: watch::Receiver<usize>,
    enabled_rx: watch::Receiver<bool>,
}

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

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

pub(super) struct RemoteControlChannels {
    pub(super) transport_event_tx: mpsc::Sender<TransportEvent>,
    pub(super) status_publisher: RemoteControlStatusPublisher,
}

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

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

    fn publish_status(&self, connection_status: RemoteControlConnectionStatus) {
        self.tx.send_if_modified(|status| {
            let next_status = RemoteControlStatusChangedNotification {
                status: connection_status,
                installation_id: status.installation_id.clone(),
                environment_id: if connection_status == RemoteControlConnectionStatus::Disabled {
                    None
                } else {
                    status.environment_id.clone()
                },
            };
            if *status == next_status {
                return false;
            }

            *status = next_status;
            true
        });
    }

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

            *status = next_status;
            true
        });
    }
}

#[derive(Clone, Copy)]
pub(super) struct RemoteControlConnectOptions<'a> {
    installation_id: &'a str,
    subscribe_cursor: Option<&'a str>,
    app_server_client_name: Option<&'a str>,
}

impl RemoteControlWebsocket {
    pub(crate) fn new(
        config: RemoteControlWebsocketConfig,
        state_db: Option<Arc<StateRuntime>>,
        auth_manager: Arc<AuthManager>,
        channels: RemoteControlChannels,
        shutdown_token: CancellationToken,
        enabled_rx: watch::Receiver<bool>,
    ) -> 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();

        Self {
            remote_control_url: config.remote_control_url,
            installation_id: config.installation_id,
            remote_control_target: config.remote_control_target,
            state_db,
            auth_manager,
            status_publisher: channels.status_publisher,
            shutdown_token,
            reconnect_attempt: 0,
            enrollment: None,
            auth_recovery,
            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,
            enabled_rx,
        }
    }

    #[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>>,
    ) {
        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(_) => {
                self.client_tracker.lock().await.shutdown().await;
                return;
            }
        };

        loop {
            if !self.wait_until_enabled().await {
                break;
            }

            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,
            };

            self.run_connection(websocket_connection, shutdown_token)
                .await;
        }

        self.client_tracker.lock().await.shutdown().await;
    }

    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),
        }
    }

    async fn wait_until_enabled(&mut self) -> bool {
        tokio::select! {
            _ = self.shutdown_token.cancelled() => false,
            enabled = self.enabled_rx.wait_for(|enabled| *enabled) => enabled.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.enabled_rx.wait_for(|enabled| !*enabled) => {
                            if changed.is_err() {
                                return ConnectOutcome::Shutdown;
                            }
                            return ConnectOutcome::Disabled;
                        }
                    }
                }
            },
        };

        loop {
            let subscribe_cursor = self.state.lock().await.subscribe_cursor.clone();
            let connect_options = RemoteControlConnectOptions {
                installation_id: &self.installation_id,
                subscribe_cursor: subscribe_cursor.as_deref(),
                app_server_client_name,
            };
            let connect_result = tokio::select! {
                _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown,
                changed = self.enabled_rx.wait_for(|enabled| !*enabled) => {
                    if changed.is_err() {
                        return ConnectOutcome::Shutdown;
                    }
                    return ConnectOutcome::Disabled;
                }
                connect_result = connect_remote_control_websocket(
                    &remote_control_target,
                    self.state_db.as_deref(),
                    &self.auth_manager,
                    &mut self.auth_recovery,
                    &mut self.enrollment,
                    connect_options,
                    &self.status_publisher,
                ) => connect_result,
            };

            match connect_result {
                Ok((websocket_connection, response)) => {
                    if !*self.enabled_rx.borrow() {
                        return ConnectOutcome::Disabled;
                    }
                    self.reconnect_attempt = 0;
                    self.auth_recovery = self.auth_manager.unauthorized_recovery();
                    self.status_publisher
                        .publish_status(RemoteControlConnectionStatus::Connected);
                    info!(
                        "connected to app-server remote control websocket: {}, {}",
                        remote_control_target.websocket_url,
                        format_headers(response.headers())
                    );
                    return ConnectOutcome::Connected(Box::new(websocket_connection));
                }
                Err(err) => {
                    if !*self.enabled_rx.borrow() {
                        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);
                        warn!(
                            "failed to connect to app-server remote control websocket: {}, err: {}",
                            remote_control_target.websocket_url, err
                        );
                        let reconnect_delay = backoff(self.reconnect_attempt);
                        self.reconnect_attempt += 1;
                        reconnect_delay
                    };
                    tokio::select! {
                        _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown,
                        changed = self.enabled_rx.wait_for(|enabled| !*enabled) => {
                            if changed.is_err() {
                                return ConnectOutcome::Shutdown;
                            }
                            return ConnectOutcome::Disabled;
                        }
                        _ = tokio::time::sleep(reconnect_delay) => {}
                    }
                }
            }
        }
    }

    async fn run_connection(
        &self,
        websocket_connection: WebSocketStream<MaybeTlsStream<TcpStream>>,
        shutdown_token: CancellationToken,
    ) {
        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 enabled_rx = self.enabled_rx.clone();
        tokio::select! {
            _ = shutdown_token.cancelled() => {}
            changed = enabled_rx.wait_for(|enabled| !*enabled) => {
                if changed.is_ok() {
                    self.status_publisher
                        .publish_status(RemoteControlConnectionStatus::Disabled);
                }
            }
            _ = join_set.join_next() => {}
        };
        shutdown_token.cancel();

        join_set.join_all().await;
    }

    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() => {
                    if let Err(err) = websocket_writer
                        .send(tungstenite::Message::Ping(Vec::new().into()))
                        .await
                    {
                        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 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 mut websocket_state = state.lock().await;
                if let Some(cursor) = client_envelope.cursor.as_deref() {
                    websocket_state.subscribe_cursor = Some(cursor.to_string());
                }
                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()
                {
                    websocket_state.outbound_buffer.ack(
                        &client_envelope.client_id,
                        stream_id,
                        acked_seq_id,
                        *segment_id,
                    );
                }
            }

            let closed_client =
                matches!(&client_envelope.event, ClientEvent::ClientClosed).then(|| {
                    (
                        client_envelope.client_id.clone(),
                        client_envelope.stream_id.clone(),
                    )
                });
            if client_tracker
                .handle_message(client_envelope)
                .await
                .is_err()
            {
                return Ok(());
            }
            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,
    auth: &RemoteControlConnectionAuth,
    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,
    )?;
    let mut auth_headers = tungstenite::http::HeaderMap::new();
    auth.auth_provider.add_auth_headers(&mut auth_headers);
    headers.extend(auth_headers);
    set_remote_control_header(headers, REMOTE_CONTROL_ACCOUNT_ID_HEADER, &auth.account_id)?;
    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)
}

pub(crate) async fn load_remote_control_auth(
    auth_manager: &Arc<AuthManager>,
) -> io::Result<RemoteControlConnectionAuth> {
    let mut reloaded = false;
    let auth = loop {
        let Some(auth) = auth_manager.auth().await else {
            if reloaded {
                return Err(io::Error::new(
                    ErrorKind::PermissionDenied,
                    "remote control requires ChatGPT authentication",
                ));
            }
            auth_manager.reload().await;
            reloaded = true;
            continue;
        };
        if !auth.uses_codex_backend() {
            break auth;
        }
        if auth.get_account_id().is_none() && !reloaded {
            auth_manager.reload().await;
            reloaded = true;
            continue;
        }
        break auth;
    };

    if !auth.uses_codex_backend() {
        return Err(io::Error::new(
            ErrorKind::PermissionDenied,
            "remote control requires ChatGPT authentication; API key auth is not supported",
        ));
    }

    Ok(RemoteControlConnectionAuth {
        auth_provider: codex_model_provider::auth_provider_from_auth(&auth),
        account_id: auth.get_account_id().ok_or_else(|| {
            io::Error::new(
                ErrorKind::WouldBlock,
                "remote control enrollment is waiting for a ChatGPT account id",
            )
        })?,
    })
}

pub(super) async fn connect_remote_control_websocket(
    remote_control_target: &RemoteControlTarget,
    state_db: Option<&StateRuntime>,
    auth_manager: &Arc<AuthManager>,
    auth_recovery: &mut UnauthorizedRecovery,
    enrollment: &mut Option<RemoteControlEnrollment>,
    connect_options: RemoteControlConnectOptions<'_>,
    status_publisher: &RemoteControlStatusPublisher,
) -> io::Result<(
    WebSocketStream<MaybeTlsStream<TcpStream>>,
    tungstenite::http::Response<()>,
)> {
    ensure_rustls_crypto_provider();

    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_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) {
        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 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;
    }

    if enrollment.is_none() {
        info!(
            "creating new remote control enrollment: websocket_url={}, enroll_url={}, account_id={}",
            remote_control_target.websocket_url, remote_control_target.enroll_url, auth.account_id
        );
        let new_enrollment = match enroll_remote_control_server(
            remote_control_target,
            &auth,
            connect_options.installation_id,
        )
        .await
        {
            Ok(new_enrollment) => new_enrollment,
            Err(err)
                if err.kind() == ErrorKind::PermissionDenied
                    && recover_remote_control_auth(auth_recovery).await =>
            {
                return Err(io::Error::other(format!(
                    "{err}; retrying after auth recovery"
                )));
            }
            Err(err) => return Err(err),
        };
        if let Err(err) = update_persisted_remote_control_enrollment(
            Some(state_db),
            remote_control_target,
            &auth.account_id,
            connect_options.app_server_client_name,
            Some(&new_enrollment),
        )
        .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);
    }

    let enrollment_ref = enrollment.as_ref().ok_or_else(|| {
        io::Error::other("missing remote control enrollment after enrollment step")
    })?;
    let request = build_remote_control_websocket_request(
        &remote_control_target.websocket_url,
        enrollment_ref,
        &auth,
        connect_options.installation_id,
        connect_options.subscribe_cursor,
    )?;

    match connect_async(request).await {
        Ok((websocket_stream, response)) => Ok((websocket_stream, response.map(|_| ()))),
        Err(err) => {
            match &err {
                tungstenite::Error::Http(response) if response.status().as_u16() == 404 => {
                    info!(
                        "remote control websocket returned HTTP 404; clearing stale enrollment before re-enrolling: websocket_url={}, account_id={}, server_id={}, environment_id={}",
                        remote_control_target.websocket_url,
                        auth.account_id,
                        enrollment_ref.server_id,
                        enrollment_ref.environment_id
                    );
                    if let Err(clear_err) = update_persisted_remote_control_enrollment(
                        Some(state_db),
                        remote_control_target,
                        &auth.account_id,
                        connect_options.app_server_client_name,
                        /*enrollment*/ None,
                    )
                    .await
                    {
                        warn!(
                            "failed to clear stale remote control enrollment in sqlite state db: {clear_err}"
                        );
                    }
                    *enrollment = None;
                    status_publisher.publish_environment_id(/*environment_id*/ None);
                }
                tungstenite::Error::Http(response)
                    if matches!(response.status().as_u16(), 401 | 403) =>
                {
                    if recover_remote_control_auth(auth_recovery).await {
                        return Err(io::Error::other(format!(
                            "remote control websocket auth failed with HTTP {}; retrying after auth recovery",
                            response.status()
                        )));
                    }
                }
                _ => {}
            }
            Err(io::Error::other(
                format_remote_control_websocket_connect_error(
                    &remote_control_target.websocket_url,
                    &err,
                ),
            ))
        }
    }
}

async fn recover_remote_control_auth(auth_recovery: &mut UnauthorizedRecovery) -> bool {
    if !auth_recovery.has_next() {
        return false;
    }

    let mode = auth_recovery.mode_name();
    let step = auth_recovery.step_name();
    match auth_recovery.next().await {
        Ok(step_result) => {
            info!(
                "remote control websocket auth recovery succeeded: mode={mode}, step={step}, auth_state_changed={:?}",
                step_result.auth_state_changed()
            );
            true
        }
        Err(err) => {
            warn!("remote control websocket auth recovery failed: mode={mode}, step={step}: {err}");
            false
        }
    }
}

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::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::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";

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

    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,
        }
    }

    #[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 enrollment = Some(RemoteControlEnrollment {
            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(),
        });
        let (status_publisher, status_rx) = remote_control_status_channel();

        let err = match connect_remote_control_websocket(
            &remote_control_target,
            Some(state_db.as_ref()),
            &auth_manager,
            &mut auth_recovery,
            &mut enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                subscribe_cursor: None,
                app_server_client_name: None,
            },
            &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_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Connecting,
                installation_id: TEST_INSTALLATION_ID.to_string(),
                environment_id: Some("env_test".to_string()),
            }
        );
    }

    #[tokio::test]
    async fn connect_remote_control_websocket_recovers_after_unauthorized_reload() {
        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");
        save_auth(
            codex_home.path(),
            &remote_control_auth_dot_json("stale-token"),
            AuthCredentialsStoreMode::File,
        )
        .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,
        )
        .await;
        let mut auth_recovery = auth_manager.unauthorized_recovery();
        let mut enrollment = Some(RemoteControlEnrollment {
            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(),
        });
        let (status_publisher, status_rx) = remote_control_status_channel();
        save_auth(
            codex_home.path(),
            &remote_control_auth_dot_json("fresh-token"),
            AuthCredentialsStoreMode::File,
        )
        .expect("fresh auth should save");

        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()),
            &auth_manager,
            &mut auth_recovery,
            &mut enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                subscribe_cursor: None,
                app_server_client_name: None,
            },
            &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,
                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; retrying after auth recovery"
        );
        assert_eq!(
            auth_manager
                .auth()
                .await
                .expect("auth should remain available")
                .get_token()
                .expect("token should be readable"),
            "fresh-token"
        );
    }

    #[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,
        )
        .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,
        )
        .await;
        let mut auth_recovery = auth_manager.unauthorized_recovery();
        let mut 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,
        )
        .expect("fresh auth should save");

        let err = connect_remote_control_websocket(
            &remote_control_target,
            Some(state_db.as_ref()),
            &auth_manager,
            &mut auth_recovery,
            &mut enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                subscribe_cursor: None,
                app_server_client_name: None,
            },
            &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"
        );
    }

    #[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 enrollment = Some(RemoteControlEnrollment {
            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(),
        });
        let (status_publisher, _status_rx) = remote_control_status_channel();

        let err = connect_remote_control_websocket(
            &remote_control_target,
            /*state_db*/ None,
            &auth_manager,
            &mut auth_recovery,
            &mut enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                subscribe_cursor: None,
                app_server_client_name: None,
            },
            &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!(enrollment, 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,
        )
        .await;
        let mut auth_recovery = auth_manager.unauthorized_recovery();
        let mut enrollment = Some(RemoteControlEnrollment {
            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(),
        });
        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()),
            &auth_manager,
            &mut auth_recovery,
            &mut enrollment,
            RemoteControlConnectOptions {
                installation_id: TEST_INSTALLATION_ID,
                subscribe_cursor: None,
                app_server_client_name: None,
            },
            &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!(enrollment, None);
        assert_eq!(
            status_rx.borrow().clone(),
            RemoteControlStatusChangedNotification {
                status: RemoteControlConnectionStatus::Connecting,
                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 (_enabled_tx, enabled_rx) = watch::channel(true);
        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),
                    },
                    /*state_db*/ None,
                    remote_control_auth_manager(),
                    RemoteControlChannels {
                        transport_event_tx,
                        status_publisher,
                    },
                    shutdown_token,
                    enabled_rx,
                )
                .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,
                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,
                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,
                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,
                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 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
        ));
        assert!(matches!(
            observe_client_message(&mut state, second_chunk),
            ClientSegmentObservation::Forward(_)
        ));
        assert!(matches!(
            observe_client_message(&mut state, first_chunk),
            ClientSegmentObservation::Dropped
        ));
    }

    #[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");
    }
}
