use super::CurrentRemoteControlEnrollment;use super::RemoteControlEnrollmentSelection;use super::RemoteControlPairingPersistenceKey;use super::desired_state::RemoteControlDesiredState;use super::desired_state::acquire_persistence_lock;use super::desired_state::desired_state_from_persisted_enrollment;use super::protocol::ClientEnvelope;use super::protocol::ClientEvent;use super::protocol::ClientId;use super::protocol::RemoteControlTarget;use super::protocol::ServerEnvelope;use super::protocol::StreamId;use super::remote_control_status_with_connection_status;use super::same_remote_control_enrollment;use super::segment::ClientSegmentObservation;use super::segment::ClientSegmentReassembler;use super::segment::REMOTE_CONTROL_SEGMENT_MAX_BYTES;use super::segment::split_server_envelope_for_transport;use crate::transport::TransportEvent;use crate::transport::remote_control::auth::RemoteControlConnectionAuth;use crate::transport::remote_control::auth::load_remote_control_auth;use crate::transport::remote_control::auth::recover_remote_control_auth;use crate::transport::remote_control::client_tracker::ClientTracker;use crate::transport::remote_control::client_tracker::REMOTE_CONTROL_IDLE_SWEEP_INTERVAL;use crate::transport::remote_control::enroll::RemoteControlEnrollment;use crate::transport::remote_control::enroll::enroll_remote_control_server;use crate::transport::remote_control::enroll::format_headers;use crate::transport::remote_control::enroll::load_persisted_remote_control_enrollment;use crate::transport::remote_control::enroll::preview_remote_control_response_body;use crate::transport::remote_control::enroll::refresh_remote_control_server;use crate::transport::remote_control::enroll::update_persisted_remote_control_enrollment;use axum::http::HeaderValue;use base64::Engine;use codex_app_server_protocol::RemoteControlConnectionStatus;use codex_app_server_protocol::RemoteControlStatusChangedNotification;use codex_core::util::backoff;use codex_login::AuthManager;use codex_login::UnauthorizedRecovery;use codex_state::StateRuntime;use codex_utils_rustls_provider::ensure_rustls_crypto_provider;use futures::SinkExt;use futures::StreamExt;use futures::stream::SplitSink;use futures::stream::SplitStream;use std::collections::HashMap;use std::collections::VecDeque;use std::io;use std::io::ErrorKind;use std::sync::Arc;use tokio::net::TcpStream;use tokio::sync::Mutex;use tokio::sync::Semaphore;use tokio::sync::mpsc;use tokio::sync::oneshot;use tokio::sync::watch;use tokio::time::MissedTickBehavior;use tokio_tungstenite::MaybeTlsStream;use tokio_tungstenite::WebSocketStream;use tokio_tungstenite::connect_async;use tokio_tungstenite::tungstenite;use tokio_tungstenite::tungstenite::client::IntoClientRequest;use tokio_util::sync::CancellationToken;#[cfg(test)]use super::RemoteControlEnrollmentState;use tracing::error;use tracing::info;use tracing::warn;pub(super) const REMOTE_CONTROL_PROTOCOL_VERSION: &str = "3";pub(super) const REMOTE_CONTROL_INSTALLATION_ID_HEADER: &str = "x-codex-installation-id";const REMOTE_CONTROL_SUBSCRIBE_CURSOR_HEADER: &str = "x-codex-subscribe-cursor";const REMOTE_CONTROL_WEBSOCKET_PING_INTERVAL: std::time::Duration = std::time::Duration::from_secs(10);const REMOTE_CONTROL_WEBSOCKET_PONG_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(60);const REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL: std::time::Duration = std::time::Duration::from_secs(1);const REMOTE_CONTROL_RECONNECT_BACKOFF_CAP: std::time::Duration = std::time::Duration::from_secs(30);const REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(30);const REMOTE_CONTROL_CONNECTION_SHUTDOWN_TIMEOUT: std::time::Duration = std::time::Duration::from_secs(5);const REMOTE_APP_SERVER_NOT_FOUND_DETAIL: &str = "Remote app server not found";struct BoundedOutboundBuffer { buffer_by_stream: HashMap<(ClientId, StreamId), VecDeque<ServerEnvelope>>, used_tx: watch::Sender<usize>,}impl BoundedOutboundBuffer { fn new() -> (Self, watch::Receiver<usize>) { let (used_tx, used_rx) = watch::channel(0); let buffer = Self { buffer_by_stream: HashMap::new(), used_tx, }; (buffer, used_rx) } fn insert(&mut self, server_envelope: &ServerEnvelope) { self.buffer_by_stream .entry(( server_envelope.client_id.clone(), server_envelope.stream_id.clone(), )) .or_default() .push_back(server_envelope.clone()); self.used_tx.send_modify(|used| *used += 1); } fn ack( &mut self, client_id: &ClientId, stream_id: &StreamId, acked_seq_id: u64, acked_segment_id: Option<usize>, ) { let key = (client_id.clone(), stream_id.clone()); let Some(buffer) = self.buffer_by_stream.get_mut(&key) else { return; }; let acked_cursor = (acked_seq_id, acked_segment_id.unwrap_or(usize::MAX)); buffer.retain(|server_envelope| { let envelope_cursor = ( server_envelope.seq_id, server_envelope.event.segment_id().unwrap_or_default(), ); let is_acked = envelope_cursor <= acked_cursor; if is_acked { self.used_tx.send_modify(|used| *used -= 1); } !is_acked }); if buffer.is_empty() { self.buffer_by_stream.remove(&key); } } fn server_envelopes(&self) -> impl Iterator<Item = &ServerEnvelope> { self.buffer_by_stream .values() .flat_map(|buffer| buffer.iter()) }}struct WebsocketState { outbound_buffer: BoundedOutboundBuffer, subscribe_cursor: Option<String>, next_seq_id_by_stream: HashMap<(ClientId, StreamId), u64>, last_completed_client_chunk_seq_id_by_stream: HashMap<(ClientId, Option<StreamId>), u64>, client_segment_reassembler: ClientSegmentReassembler,}impl WebsocketState { fn observe_client_message( &mut self, client_envelope: ClientEnvelope, wire_size_bytes: usize, ) -> ClientSegmentObservation { let client_message_key = Self::client_message_key(&client_envelope); if let Some((key, seq_id)) = client_message_key.as_ref() && self .last_completed_client_chunk_seq_id_by_stream .get(key) .is_some_and(|last_seq_id| last_seq_id >= seq_id) { return ClientSegmentObservation::Dropped; } if let ( Some((_, seq_id)), Some(stream_id), ClientEvent::ClientMessageChunk { segment_id, .. }, ) = ( client_message_key.as_ref(), client_envelope.stream_id.as_ref(), &client_envelope.event, ) && self.client_segment_reassembler.should_ignore_chunk( &client_envelope.client_id, stream_id, *seq_id, *segment_id, ) { return ClientSegmentObservation::Dropped; } if client_message_key.is_some() && wire_size_bytes > REMOTE_CONTROL_SEGMENT_MAX_BYTES { warn!( client_id = client_envelope.client_id.0.as_str(), "dropping oversized segmented remote-control client envelope" ); if let Some(stream_id) = client_envelope.stream_id.as_ref() { self.client_segment_reassembler .invalidate_stream(&client_envelope.client_id, stream_id); } return ClientSegmentObservation::Dropped; } self.client_segment_reassembler.observe(client_envelope) } fn record_client_message_delivery( &mut self, client_envelope: &ClientEnvelope, client_message_key: Option<((ClientId, Option<StreamId>), u64)>, ) { if let Some(cursor) = client_envelope.cursor.as_deref() { self.subscribe_cursor = Some(cursor.to_string()); } if let Some((key, seq_id)) = client_message_key { self.last_completed_client_chunk_seq_id_by_stream .insert(key, seq_id); } if let ClientEvent::Ack { segment_id } = &client_envelope.event && let Some(acked_seq_id) = client_envelope.seq_id && let Some(stream_id) = client_envelope.stream_id.as_ref() { self.outbound_buffer.ack( &client_envelope.client_id, stream_id, acked_seq_id, *segment_id, ); } } fn invalidate_client_message_stream(&mut self, client_id: &ClientId, stream_id: &StreamId) { self.last_completed_client_chunk_seq_id_by_stream .remove(&(client_id.clone(), Some(stream_id.clone()))); } fn invalidate_client_message_client(&mut self, client_id: &ClientId) { self.last_completed_client_chunk_seq_id_by_stream .retain(|(cursor_client_id, _), _| cursor_client_id != client_id); } fn client_message_key( client_envelope: &ClientEnvelope, ) -> Option<((ClientId, Option<StreamId>), u64)> { let seq_id = match (&client_envelope.event, client_envelope.seq_id) { (ClientEvent::ClientMessageChunk { .. }, Some(seq_id)) => seq_id, _ => return None, }; Some(( ( client_envelope.client_id.clone(), client_envelope.stream_id.clone(), ), seq_id, )) }}pub(crate) struct RemoteControlWebsocket { remote_control_url: String, installation_id: String, server_name: String, remote_control_target: Option<RemoteControlTarget>, state_db: Option<Arc<StateRuntime>>, auth_manager: Arc<AuthManager>, status_publisher: RemoteControlStatusPublisher, shutdown_token: CancellationToken, reconnect_attempt: u64, auth_recovery: UnauthorizedRecovery, auth_change_rx: watch::Receiver<u64>, current_enrollment: CurrentRemoteControlEnrollment, pairing_persistence_key: RemoteControlPairingPersistenceKey, client_tracker: Arc<Mutex<ClientTracker>>, state: Arc<Mutex<WebsocketState>>, server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>, used_rx: watch::Receiver<usize>, desired_state_tx: Arc<watch::Sender<RemoteControlDesiredState>>, desired_state_rx: watch::Receiver<RemoteControlDesiredState>, desired_state_persistence_lock: Arc<Semaphore>,}pub(crate) struct RemoteControlWebsocketConfig { pub(crate) remote_control_url: String, pub(crate) installation_id: String, pub(crate) remote_control_target: Option<RemoteControlTarget>, pub(crate) server_name: String,}pub(super) struct RemoteControlAuthContext<'a> { auth_manager: &'a Arc<AuthManager>, auth_recovery: &'a mut UnauthorizedRecovery, auth_change_rx: &'a mut watch::Receiver<u64>,}struct RemoteControlEnrollmentAuthContext<'a, 'b> { auth: &'a RemoteControlConnectionAuth, recovery: &'a mut RemoteControlAuthContext<'b>,}enum ConnectOutcome { Connected(Box<WebSocketStream<MaybeTlsStream<TcpStream>>>), Disabled, Shutdown,}#[derive(Debug, Clone, Copy)]enum ConnectionEndReason { Shutdown, Disabled, EnabledWatchClosed, ConnectionWorkerStopped,}pub(super) struct RemoteControlChannels { pub(super) transport_event_tx: mpsc::Sender<TransportEvent>, pub(super) status_publisher: RemoteControlStatusPublisher, pub(super) current_enrollment: CurrentRemoteControlEnrollment, pub(super) pairing_persistence_key: RemoteControlPairingPersistenceKey, pub(super) desired_state_persistence_lock: Arc<Semaphore>,}#[derive(Clone)]pub(super) struct RemoteControlStatusPublisher { tx: watch::Sender<RemoteControlStatusChangedNotification>,}impl RemoteControlStatusPublisher { pub(super) fn new(tx: watch::Sender<RemoteControlStatusChangedNotification>) -> Self { Self { tx } } fn status(&self) -> RemoteControlStatusChangedNotification { self.tx.borrow().clone() } fn publish_status(&self, connection_status: RemoteControlConnectionStatus) { let mut status_change = None; self.tx.send_if_modified(|status| { let next_status = remote_control_status_with_connection_status(status, connection_status); if *status == next_status { return false; } status_change = Some((status.clone(), next_status.clone())); *status = next_status; true }); if let Some((previous_status, next_status)) = status_change { info!( previous_status = ?previous_status.status, next_status = ?next_status.status, previous_environment_id = ?previous_status.environment_id, next_environment_id = ?next_status.environment_id, installation_id = %next_status.installation_id, server_name = %next_status.server_name, "remote control websocket status changed" ); } } pub(super) fn publish_environment_id(&self, environment_id: Option<String>) { let mut status_change = None; self.tx.send_if_modified(|status| { if status.status == RemoteControlConnectionStatus::Disabled { return false; } let next_status = RemoteControlStatusChangedNotification { status: status.status, server_name: status.server_name.clone(), installation_id: status.installation_id.clone(), environment_id, }; if *status == next_status { return false; } status_change = Some((status.clone(), next_status.clone())); *status = next_status; true }); if let Some((previous_status, next_status)) = status_change { info!( status = ?next_status.status, previous_environment_id = ?previous_status.environment_id, next_environment_id = ?next_status.environment_id, installation_id = %next_status.installation_id, server_name = %next_status.server_name, "remote control websocket environment changed" ); } }}#[derive(Clone, Copy)]pub(super) struct RemoteControlConnectOptions<'a> { installation_id: &'a str, server_name: &'a str, subscribe_cursor: Option<&'a str>, app_server_client_name: Option<&'a str>, desired_state_tx: &'a watch::Sender<RemoteControlDesiredState>, desired_state_persistence_lock: &'a Semaphore,}impl RemoteControlWebsocket { pub(crate) fn new( config: RemoteControlWebsocketConfig, state_db: Option<Arc<StateRuntime>>, auth_manager: Arc<AuthManager>, channels: RemoteControlChannels, shutdown_token: CancellationToken, desired_state_tx: Arc<watch::Sender<RemoteControlDesiredState>>, ) -> Self { let shutdown_token = shutdown_token.child_token(); let (server_event_tx, server_event_rx) = mpsc::channel(super::CHANNEL_CAPACITY); let client_tracker = ClientTracker::new( server_event_tx, channels.transport_event_tx, &shutdown_token, ); let (outbound_buffer, used_rx) = BoundedOutboundBuffer::new(); let auth_recovery = auth_manager.unauthorized_recovery(); let auth_change_rx = auth_manager.auth_change_receiver(); let desired_state_rx = desired_state_tx.subscribe(); Self { remote_control_url: config.remote_control_url, installation_id: config.installation_id, server_name: config.server_name, remote_control_target: config.remote_control_target, state_db, auth_manager, status_publisher: channels.status_publisher, shutdown_token, reconnect_attempt: 0, auth_recovery, auth_change_rx, current_enrollment: channels.current_enrollment, pairing_persistence_key: channels.pairing_persistence_key, client_tracker: Arc::new(Mutex::new(client_tracker)), state: Arc::new(Mutex::new(WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), })), server_event_rx: Arc::new(Mutex::new(server_event_rx)), used_rx, desired_state_tx, desired_state_rx, desired_state_persistence_lock: channels.desired_state_persistence_lock, } } #[expect( clippy::await_holding_invalid_type, reason = "remote-control client shutdown must serialize tracker state" )] pub(crate) async fn run( mut self, app_server_client_name_rx: Option<oneshot::Receiver<String>>, ) { info!( remote_control_url = %self.remote_control_url, installation_id = %self.installation_id, server_name = %self.server_name, initial_desired_state = ?*self.desired_state_rx.borrow(), "app-server remote control websocket loop started" ); let app_server_client_name = match self .wait_for_app_server_client_name(app_server_client_name_rx) .await { Ok(app_server_client_name) => app_server_client_name, Err(_) => { warn!( remote_control_url = %self.remote_control_url, installation_id = %self.installation_id, server_name = %self.server_name, shutdown_requested = self.shutdown_token.is_cancelled(), "app-server remote control websocket loop stopped before client name was ready" ); self.client_tracker.lock().await.shutdown().await; return; } }; self.pairing_persistence_key .send_replace(app_server_client_name.clone()); if matches!( *self.desired_state_rx.borrow(), RemoteControlDesiredState::Unknown ) && !self .resolve_unknown_desired_state(app_server_client_name.as_deref()) .await { self.client_tracker.lock().await.shutdown().await; return; } loop { if !self.wait_until_enabled().await { info!( remote_control_url = %self.remote_control_url, installation_id = %self.installation_id, server_name = %self.server_name, shutdown_requested = self.shutdown_token.is_cancelled(), current_status = ?self.status_publisher.status().status, "app-server remote control websocket loop exiting while waiting for enablement" ); break; } let status = self.status_publisher.status(); info!( remote_control_url = %self.remote_control_url, installation_id = %self.installation_id, server_name = %self.server_name, reconnect_attempt = self.reconnect_attempt.saturating_add(1), current_status = ?status.status, environment_id = ?status.environment_id, "starting app-server remote control websocket connection cycle" ); let shutdown_token = self.shutdown_token.child_token(); let websocket_connection = match self .connect(&shutdown_token, app_server_client_name.as_deref()) .await { ConnectOutcome::Connected(websocket_connection) => *websocket_connection, ConnectOutcome::Disabled => { self.status_publisher .publish_status(RemoteControlConnectionStatus::Disabled); continue; } ConnectOutcome::Shutdown => break, }; let connection_end_reason = self .run_connection(websocket_connection, shutdown_token) .await; let status = self.status_publisher.status(); info!( remote_control_url = %self.remote_control_url, installation_id = %self.installation_id, server_name = %self.server_name, connection_end_reason = ?connection_end_reason, current_status = ?status.status, environment_id = ?status.environment_id, desired_state = ?*self.desired_state_rx.borrow(), "app-server remote control websocket connection cycle ended" ); } self.client_tracker.lock().await.shutdown().await; info!( remote_control_url = %self.remote_control_url, installation_id = %self.installation_id, server_name = %self.server_name, shutdown_requested = self.shutdown_token.is_cancelled(), "app-server remote control websocket loop exited" ); } async fn wait_for_app_server_client_name( &self, app_server_client_name_rx: Option<oneshot::Receiver<String>>, ) -> Result<Option<String>, ()> { match app_server_client_name_rx { Some(app_server_client_name_rx) => { tokio::select! { _ = self.shutdown_token.cancelled() => Err(()), app_server_client_name = app_server_client_name_rx => match app_server_client_name { Ok(app_server_client_name) => Ok(Some(app_server_client_name)), Err(_) => Err(()), }, } } None => Ok(None), } } pub(super) async fn resolve_unknown_desired_state( &mut self, app_server_client_name: Option<&str>, ) -> bool { let remote_control_target = match super::protocol::normalize_remote_control_url( &self.remote_control_url, ) { Ok(remote_control_target) => remote_control_target, Err(err) => { warn!( "remote control preference cannot be resolved because the URL is invalid: {err}" ); self.transition_unknown_to(RemoteControlDesiredState::Disabled); return true; } }; self.remote_control_target = Some(remote_control_target.clone()); let Some(state_db) = self.state_db.clone() else { self.transition_unknown_to(RemoteControlDesiredState::Disabled); return true; }; loop { if !matches!( *self.desired_state_rx.borrow(), RemoteControlDesiredState::Unknown ) { return true; } let auth = match load_remote_control_auth(&self.auth_manager).await { Ok(auth) => auth, Err(err) => { info!( error = %err, "waiting to resolve remote control preference until authentication is available" ); if !self.wait_for_preference_resolution_retry().await { return false; } continue; } }; let enrollment = match state_db .get_remote_control_enrollment( &remote_control_target.websocket_url, &auth.account_id, app_server_client_name, ) .await { Ok(enrollment) => enrollment, Err(err) => { warn!( error = %err, "failed to resolve persisted remote control preference; retrying" ); if !self.wait_for_preference_resolution_retry().await { return false; } continue; } }; let desired_state = desired_state_from_persisted_enrollment(enrollment); self.transition_unknown_to(desired_state); return true; } } fn transition_unknown_to(&self, desired_state: RemoteControlDesiredState) { self.desired_state_tx.send_if_modified(|state| { if !matches!(*state, RemoteControlDesiredState::Unknown) { return false; } *state = desired_state; true }); } async fn wait_for_preference_resolution_retry(&mut self) -> bool { tokio::select! { _ = self.shutdown_token.cancelled() => false, changed = self.desired_state_rx.changed() => changed.is_ok(), _ = tokio::time::sleep(REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL) => true, } } async fn wait_until_enabled(&mut self) -> bool { tokio::select! { _ = self.shutdown_token.cancelled() => false, desired_state = self.desired_state_rx.wait_for(|state| state.is_enabled()) => desired_state.is_ok(), } } async fn connect( &mut self, shutdown_token: &CancellationToken, app_server_client_name: Option<&str>, ) -> ConnectOutcome { self.status_publisher .publish_status(RemoteControlConnectionStatus::Connecting); let remote_control_target = match self.remote_control_target.as_ref() { Some(remote_control_target) => remote_control_target.clone(), None => match super::protocol::normalize_remote_control_url(&self.remote_control_url) { Ok(remote_control_target) => { self.remote_control_target = Some(remote_control_target.clone()); remote_control_target } Err(err) => { self.status_publisher .publish_status(RemoteControlConnectionStatus::Errored); warn!("remote control is enabled but the URL is invalid: {err}"); tokio::select! { _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown, changed = self.desired_state_rx.wait_for(|state| !state.is_enabled()) => { if changed.is_err() { return ConnectOutcome::Shutdown; } return ConnectOutcome::Disabled; } } } }, }; loop { let subscribe_cursor = self.state.lock().await.subscribe_cursor.clone(); let enrollment = self.current_enrollment.snapshot(); info!( websocket_url = %remote_control_target.websocket_url, installation_id = %self.installation_id, server_name = %self.server_name, reconnect_attempt = self.reconnect_attempt.saturating_add(1), has_enrollment = enrollment.is_some(), server_id = ?enrollment.as_ref().map(|enrollment| enrollment.server_id.as_str()), environment_id = ?enrollment.as_ref().map(|enrollment| enrollment.environment_id.as_str()), subscribe_cursor_present = subscribe_cursor.is_some(), app_server_client_name = ?app_server_client_name, "connecting to app-server remote control websocket" ); let connect_options = RemoteControlConnectOptions { installation_id: &self.installation_id, server_name: &self.server_name, subscribe_cursor: subscribe_cursor.as_deref(), app_server_client_name, desired_state_tx: &self.desired_state_tx, desired_state_persistence_lock: &self.desired_state_persistence_lock, }; let auth_context = RemoteControlAuthContext { auth_manager: &self.auth_manager, auth_recovery: &mut self.auth_recovery, auth_change_rx: &mut self.auth_change_rx, }; let mut disabled_rx = self.desired_state_rx.clone(); let connect_result = tokio::select! { _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown, changed = disabled_rx.wait_for(|state| !state.is_enabled()) => { if changed.is_err() { return ConnectOutcome::Shutdown; } return ConnectOutcome::Disabled; } connect_result = async { connect_remote_control_websocket( &remote_control_target, self.state_db.as_deref(), auth_context, &self.current_enrollment, connect_options, &self.status_publisher, ) .await } => connect_result, }; match connect_result { Ok((websocket_connection, response)) => { if !self.desired_state_rx.borrow().is_enabled() { return ConnectOutcome::Disabled; } self.reconnect_attempt = 0; self.auth_recovery = self.auth_manager.unauthorized_recovery(); self.status_publisher .publish_status(RemoteControlConnectionStatus::Connected); let enrollment = self.current_enrollment.snapshot(); info!( websocket_url = %remote_control_target.websocket_url, installation_id = %self.installation_id, server_name = %self.server_name, server_id = ?enrollment.as_ref().map(|enrollment| enrollment.server_id.as_str()), environment_id = ?enrollment.as_ref().map(|enrollment| enrollment.environment_id.as_str()), subscribe_cursor_present = subscribe_cursor.is_some(), response_headers = %format_headers(response.headers()), "connected to app-server remote control websocket" ); return ConnectOutcome::Connected(Box::new(websocket_connection)); } Err(err) => { if !self.desired_state_rx.borrow().is_enabled() { return ConnectOutcome::Disabled; } let reconnect_delay = if err.kind() == ErrorKind::WouldBlock { REMOTE_CONTROL_ACCOUNT_ID_RETRY_INTERVAL } else { self.status_publisher .publish_status(RemoteControlConnectionStatus::Errored); let reconnect_attempt = self.reconnect_attempt.saturating_add(1); let (reconnect_delay, reconnect_backoff_reset) = next_reconnect_delay(&mut self.reconnect_attempt); let enrollment = self.current_enrollment.snapshot(); warn!( websocket_url = %remote_control_target.websocket_url, installation_id = %self.installation_id, server_name = %self.server_name, error = %err, error_kind = ?err.kind(), reconnect_attempt, reconnect_delay = ?reconnect_delay, reconnect_backoff_reset, has_enrollment = enrollment.is_some(), server_id = ?enrollment.as_ref().map(|enrollment| enrollment.server_id.as_str()), environment_id = ?enrollment.as_ref().map(|enrollment| enrollment.environment_id.as_str()), subscribe_cursor_present = subscribe_cursor.is_some(), "failed to connect to app-server remote control websocket" ); if reconnect_backoff_reset { info!( reconnect_backoff_cap = ?REMOTE_CONTROL_RECONNECT_BACKOFF_CAP, "reset app-server remote control websocket reconnect backoff after cap" ); } reconnect_delay }; tokio::select! { _ = shutdown_token.cancelled() => return ConnectOutcome::Shutdown, changed = self.desired_state_rx.wait_for(|state| !state.is_enabled()) => { if changed.is_err() { return ConnectOutcome::Shutdown; } return ConnectOutcome::Disabled; } changed = self.auth_change_rx.changed() => { if changed.is_err() { return ConnectOutcome::Shutdown; } self.auth_recovery = self.auth_manager.unauthorized_recovery(); self.reconnect_attempt = 0; info!("retrying app-server remote control websocket after auth changed"); } _ = tokio::time::sleep(reconnect_delay) => {} } } } } } async fn run_connection( &self, websocket_connection: WebSocketStream<MaybeTlsStream<TcpStream>>, shutdown_token: CancellationToken, ) -> ConnectionEndReason { let (websocket_writer, websocket_reader) = websocket_connection.split(); let mut join_set = tokio::task::JoinSet::new(); join_set.spawn(Self::run_server_writer( self.state.clone(), self.server_event_rx.clone(), self.used_rx.clone(), websocket_writer, REMOTE_CONTROL_WEBSOCKET_PING_INTERVAL, shutdown_token.clone(), )); join_set.spawn(Self::run_websocket_reader( self.client_tracker.clone(), self.state.clone(), websocket_reader, REMOTE_CONTROL_WEBSOCKET_PONG_TIMEOUT, shutdown_token.clone(), )); let mut desired_state_rx = self.desired_state_rx.clone(); let connection_end_reason = tokio::select! { _ = shutdown_token.cancelled() => ConnectionEndReason::Shutdown, changed = desired_state_rx.wait_for(|state| !state.is_enabled()) => { if changed.is_ok() { self.status_publisher .publish_status(RemoteControlConnectionStatus::Disabled); ConnectionEndReason::Disabled } else { ConnectionEndReason::EnabledWatchClosed } } _ = join_set.join_next() => ConnectionEndReason::ConnectionWorkerStopped, }; shutdown_token.cancel(); Self::join_connection_workers(&mut join_set, REMOTE_CONTROL_CONNECTION_SHUTDOWN_TIMEOUT) .await; connection_end_reason } async fn join_connection_workers( join_set: &mut tokio::task::JoinSet<()>, shutdown_timeout: std::time::Duration, ) { if tokio::time::timeout(shutdown_timeout, Self::drain_join_set(join_set)) .await .is_ok() { return; } warn!( shutdown_timeout = ?shutdown_timeout, remaining_workers = join_set.len(), "timed out waiting for remote control connection workers to stop; aborting" ); join_set.abort_all(); Self::drain_join_set(join_set).await; } async fn drain_join_set(join_set: &mut tokio::task::JoinSet<()>) { while join_set.join_next().await.is_some() {} } async fn run_server_writer( state: Arc<Mutex<WebsocketState>>, server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>, used_rx: watch::Receiver<usize>, websocket_writer: SplitSink< WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::Message, >, ping_interval: std::time::Duration, shutdown_token: CancellationToken, ) { let result = Self::run_server_writer_inner( state, server_event_rx, used_rx, websocket_writer, ping_interval, shutdown_token, ) .await; if let Err(err) = result { warn!("remote control websocket writer disconnected, err: {err}"); } else { warn!("remote control websocket writer was stopped"); } } #[expect( clippy::await_holding_invalid_type, reason = "remote-control server event receiver is shared across reconnects" )] async fn run_server_writer_inner( state: Arc<Mutex<WebsocketState>>, server_event_rx: Arc<Mutex<mpsc::Receiver<super::QueuedServerEnvelope>>>, mut used_rx: watch::Receiver<usize>, mut websocket_writer: SplitSink< WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::Message, >, ping_interval: std::time::Duration, shutdown_token: CancellationToken, ) -> io::Result<()> { let server_envelopes = state .lock() .await .outbound_buffer .server_envelopes() .cloned() .collect::<Vec<_>>(); for server_envelope in server_envelopes { let payload = match serde_json::to_string(&server_envelope) { Ok(payload) => payload, Err(err) => { error!("failed to serialize remote-control server event: {err}"); continue; } }; tokio::select! { _ = shutdown_token.cancelled() => return Ok(()), send_result = websocket_writer.send(tungstenite::Message::Text(payload.into())) => { if let Err(err) = send_result { return Err(io::Error::other(err)); } } }; } let mut ping_interval = tokio::time::interval_at(tokio::time::Instant::now() + ping_interval, ping_interval); ping_interval.set_missed_tick_behavior(MissedTickBehavior::Skip); let mut server_event_rx = server_event_rx.lock().await; loop { let outbound_has_capacity = *used_rx.borrow() < super::CHANNEL_CAPACITY; let queued_server_envelope = tokio::select! { _ = shutdown_token.cancelled() => return Ok(()), _ = ping_interval.tick() => { tokio::select! { _ = shutdown_token.cancelled() => return Ok(()), send_result = websocket_writer.send(tungstenite::Message::Ping(Vec::new().into())) => { if let Err(err) = send_result { return Err(io::Error::other(err)); } } }; continue; } wait_result = used_rx.changed(), if !outbound_has_capacity => { if wait_result.is_err() { return Err(io::Error::new( ErrorKind::UnexpectedEof, "outbound buffer usage channel closed", )); } continue; } recv_result = server_event_rx.recv(), if outbound_has_capacity => { match recv_result { Some(queued_server_envelope) => queued_server_envelope, None => { return Err(io::Error::new(io::ErrorKind::UnexpectedEof, "server event channel closed")); } } } }; let (payloads, write_complete_tx) = { let mut state = state.lock().await; let seq_key = ( queued_server_envelope.client_id.clone(), queued_server_envelope.stream_id.clone(), ); let seq_id = *state .next_seq_id_by_stream .entry(seq_key.clone()) .or_insert(1); let server_envelope = ServerEnvelope { event: queued_server_envelope.event, client_id: queued_server_envelope.client_id, seq_id, stream_id: queued_server_envelope.stream_id, }; let server_envelopes = match split_server_envelope_for_transport(server_envelope) { Ok(server_envelopes) => server_envelopes, Err(err) => { error!("failed to split remote-control server event: {err}"); continue; } }; let mut payloads = Vec::with_capacity(server_envelopes.len()); for server_envelope in server_envelopes { let payload = match serde_json::to_string(&server_envelope) { Ok(payload) => payload, Err(err) => { error!("failed to serialize remote-control server event: {err}"); continue; } }; state.outbound_buffer.insert(&server_envelope); payloads.push(payload); } state .next_seq_id_by_stream .insert(seq_key, seq_id.saturating_add(1)); (payloads, queued_server_envelope.write_complete_tx) }; for payload in payloads { tokio::select! { _ = shutdown_token.cancelled() => return Ok(()), send_result = websocket_writer.send(tungstenite::Message::Text(payload.into())) => { if let Err(err) = send_result { return Err(io::Error::other(err)); } } } } if let Some(write_complete_tx) = write_complete_tx { let _ = write_complete_tx.send(()); } } } async fn run_websocket_reader( client_tracker: Arc<Mutex<ClientTracker>>, state: Arc<Mutex<WebsocketState>>, websocket_reader: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>, pong_timeout: std::time::Duration, shutdown_token: CancellationToken, ) { let result = Self::run_websocket_reader_inner( client_tracker, state, websocket_reader, pong_timeout, shutdown_token, ) .await; if let Err(err) = result { warn!("remote control websocket reader disconnected, err: {err}"); } else { warn!("remote control websocket reader was stopped"); } } #[expect( clippy::await_holding_invalid_type, reason = "remote-control client tracking must stay serialized while processing inbound events" )] async fn run_websocket_reader_inner( client_tracker: Arc<Mutex<ClientTracker>>, state: Arc<Mutex<WebsocketState>>, mut websocket_reader: SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>, pong_timeout: std::time::Duration, shutdown_token: CancellationToken, ) -> io::Result<()> { let mut client_tracker = client_tracker.lock().await; let mut idle_sweep_interval = tokio::time::interval(REMOTE_CONTROL_IDLE_SWEEP_INTERVAL); idle_sweep_interval.set_missed_tick_behavior(MissedTickBehavior::Skip); let pong_deadline = tokio::time::sleep(pong_timeout); tokio::pin!(pong_deadline); loop { let incoming_message = tokio::select! { _ = shutdown_token.cancelled() => return Ok(()), _ = &mut pong_deadline => { return Err(io::Error::new( ErrorKind::TimedOut, "remote control websocket pong timeout", )); } client_key = client_tracker.bookkeep_join_set() => { let Some(client_key) = client_key else { continue; }; if client_tracker.close_client(&client_key).await.is_err() { return Ok(()); } state .lock() .await .client_segment_reassembler .invalidate_stream(&client_key.0, &client_key.1); state .lock() .await .invalidate_client_message_stream(&client_key.0, &client_key.1); continue; } _ = idle_sweep_interval.tick() => { match client_tracker.close_expired_clients().await { Ok(client_keys) => { let mut websocket_state = state.lock().await; for (client_id, stream_id) in client_keys { websocket_state .client_segment_reassembler .invalidate_stream(&client_id, &stream_id); websocket_state .invalidate_client_message_stream(&client_id, &stream_id); } } Err(_) => return Ok(()), } continue; } incoming_message = websocket_reader.next() => { match incoming_message { Some(incoming_message) => incoming_message, None => return Err(io::Error::new(ErrorKind::UnexpectedEof, "websocket stream ended")), } } }; let (client_envelope, wire_size_bytes) = match incoming_message { Ok(tungstenite::Message::Text(text)) => { let wire_size_bytes = text.len(); match serde_json::from_str::<ClientEnvelope>(&text) { Ok(client_envelope) => (client_envelope, wire_size_bytes), Err(err) => { warn!("failed to deserialize remote-control client event: {err}"); continue; } } } Ok(tungstenite::Message::Pong(_)) => { pong_deadline .as_mut() .reset(tokio::time::Instant::now() + pong_timeout); continue; } Ok(tungstenite::Message::Ping(_)) | Ok(tungstenite::Message::Frame(_)) => continue, Ok(tungstenite::Message::Binary(_)) => { warn!("dropping unsupported binary remote-control websocket message"); continue; } Ok(tungstenite::Message::Close(_)) => { return Err(io::Error::new( ErrorKind::ConnectionAborted, "websocket disconnected", )); } Err(err) => { return Err(io::Error::new( ErrorKind::InvalidData, format!("failed to read from websocket: {err}"), )); } }; let client_message_key = WebsocketState::client_message_key(&client_envelope); let observation = { let mut websocket_state = state.lock().await; websocket_state.observe_client_message(client_envelope, wire_size_bytes) }; let client_envelope = match observation { ClientSegmentObservation::Forward(client_envelope) => *client_envelope, ClientSegmentObservation::Pending | ClientSegmentObservation::Dropped => continue, }; let closed_client = matches!(&client_envelope.event, ClientEvent::ClientClosed).then(|| { ( client_envelope.client_id.clone(), client_envelope.stream_id.clone(), ) }); let delivered_client_envelope = client_envelope.clone(); if client_tracker .handle_message(client_envelope) .await .is_err() { return Ok(()); } state .lock() .await .record_client_message_delivery(&delivered_client_envelope, client_message_key); if let Some((client_id, stream_id)) = closed_client { let mut websocket_state = state.lock().await; if let Some(stream_id) = stream_id { websocket_state .client_segment_reassembler .invalidate_stream(&client_id, &stream_id); websocket_state.invalidate_client_message_stream(&client_id, &stream_id); } else { websocket_state .client_segment_reassembler .invalidate_client(&client_id); websocket_state.invalidate_client_message_client(&client_id); } } } }}fn set_remote_control_header( headers: &mut tungstenite::http::HeaderMap, name: &'static str, value: &str,) -> io::Result<()> { let header_value = HeaderValue::from_str(value).map_err(|err| { io::Error::new( ErrorKind::InvalidInput, format!("invalid remote control header `{name}`: {err}"), ) })?; headers.insert(name, header_value); Ok(())}fn build_remote_control_websocket_request( websocket_url: &str, enrollment: &RemoteControlEnrollment, installation_id: &str, subscribe_cursor: Option<&str>,) -> io::Result<tungstenite::http::Request<()>> { let mut request = websocket_url.into_client_request().map_err(|err| { io::Error::new( ErrorKind::InvalidInput, format!("invalid remote control websocket URL `{websocket_url}`: {err}"), ) })?; let headers = request.headers_mut(); set_remote_control_header(headers, "x-codex-server-id", &enrollment.server_id)?; set_remote_control_header( headers, "x-codex-name", &base64::engine::general_purpose::STANDARD.encode(&enrollment.server_name), )?; set_remote_control_header( headers, "x-codex-protocol-version", REMOTE_CONTROL_PROTOCOL_VERSION, )?; set_remote_control_header( headers, "authorization", &format!( "Bearer {}", enrollment .remote_control_token .as_deref() .ok_or_else(|| io::Error::other("missing remote control server token"))? ), )?; set_remote_control_header( headers, REMOTE_CONTROL_INSTALLATION_ID_HEADER, installation_id, )?; if let Some(subscribe_cursor) = subscribe_cursor { set_remote_control_header( headers, REMOTE_CONTROL_SUBSCRIBE_CURSOR_HEADER, subscribe_cursor, )?; } Ok(request)}fn next_reconnect_delay(reconnect_attempt: &mut u64) -> (std::time::Duration, bool) { let reconnect_delay = backoff(*reconnect_attempt).min(REMOTE_CONTROL_RECONNECT_BACKOFF_CAP); let reconnect_backoff_reset = reconnect_delay == REMOTE_CONTROL_RECONNECT_BACKOFF_CAP; *reconnect_attempt = if reconnect_backoff_reset { 0 } else { (*reconnect_attempt).saturating_add(1) }; (reconnect_delay, reconnect_backoff_reset)}pub(super) async fn connect_remote_control_websocket( remote_control_target: &RemoteControlTarget, state_db: Option<&StateRuntime>, mut auth_context: RemoteControlAuthContext<'_>, current_enrollment: &CurrentRemoteControlEnrollment, connect_options: RemoteControlConnectOptions<'_>, status_publisher: &RemoteControlStatusPublisher,) -> io::Result<( WebSocketStream<MaybeTlsStream<TcpStream>>, tungstenite::http::Response<()>,)> { ensure_rustls_crypto_provider(); let (auth, enrollment) = { let mut current_enrollment = current_enrollment.lock().await; let auth = prepare_remote_control_enrollment( remote_control_target, state_db, &mut auth_context, &mut current_enrollment, connect_options, status_publisher, ) .await?; let enrollment = current_enrollment.as_ref().cloned().ok_or_else(|| { io::Error::other("missing remote control enrollment after enrollment step") })?; (auth, enrollment) }; let request = build_remote_control_websocket_request( &remote_control_target.websocket_url, &enrollment, connect_options.installation_id, connect_options.subscribe_cursor, )?; let websocket_connect_result = tokio::time::timeout( REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT, connect_async(request), ) .await .map_err(|_| { io::Error::new( ErrorKind::TimedOut, format!( "timed out connecting to remote control websocket at `{}` after {:?}", remote_control_target.websocket_url, REMOTE_CONTROL_WEBSOCKET_CONNECT_TIMEOUT ), ) })?; match websocket_connect_result { Ok((websocket_stream, response)) => Ok((websocket_stream, response.map(|_| ()))), Err(err) => { match &err { tungstenite::Error::Http(response) if websocket_response_reports_missing_remote_app_server(response) => { info!( "remote control websocket returned HTTP 404; replacing stale enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}", remote_control_target.websocket_url, auth.account_id, enrollment.server_id, enrollment.environment_id ); replace_remote_control_enrollment_if_matches( state_db, remote_control_target, RemoteControlEnrollmentAuthContext { auth: &auth, recovery: &mut auth_context, }, current_enrollment, &enrollment, connect_options, status_publisher, ) .await?; } tungstenite::Error::Http(response) if response.status().as_u16() == 404 => { let response_body = response .body() .as_deref() .map(preview_remote_control_response_body) .unwrap_or_else(|| "<missing>".to_string()); warn!( websocket_url = %remote_control_target.websocket_url, account_id = %auth.account_id, server_id = %enrollment.server_id, environment_id = %enrollment.environment_id, response_status = %response.status(), response_headers = %format_headers(response.headers()), response_body = %response_body, "remote control websocket returned unrecognized HTTP 404; preserving enrollment before retry" ); } tungstenite::Error::Http(response) if matches!(response.status().as_u16(), 401 | 403) => { clear_remote_control_server_token_if_matches(current_enrollment, &enrollment) .await?; return Err(io::Error::other(format!( "remote control websocket auth failed with HTTP {}; refreshing server token before reconnect", response.status() ))); } _ => {} } Err(io::Error::other( format_remote_control_websocket_connect_error( &remote_control_target.websocket_url, &err, ), )) } }}async fn prepare_remote_control_enrollment( remote_control_target: &RemoteControlTarget, state_db: Option<&StateRuntime>, auth_context: &mut RemoteControlAuthContext<'_>, enrollment: &mut Option<RemoteControlEnrollment>, connect_options: RemoteControlConnectOptions<'_>, status_publisher: &RemoteControlStatusPublisher,) -> io::Result<RemoteControlConnectionAuth> { let Some(state_db) = state_db else { *enrollment = None; return Err(io::Error::new( ErrorKind::NotFound, "remote control requires sqlite state db", )); }; let auth = match load_remote_control_auth(auth_context.auth_manager).await { Ok(auth) => auth, Err(err) => { if err.kind() == ErrorKind::PermissionDenied { *enrollment = None; status_publisher.publish_environment_id(/*environment_id*/ None); } return Err(err); } }; let enrollment_account_id = enrollment.as_ref().map(|enrollment| &enrollment.account_id); if enrollment_account_id.is_some_and(|account_id| account_id != &auth.account_id) { resolve_desired_state_after_account_change( state_db, remote_control_target, auth_context.auth_manager, &auth.account_id, connect_options, ) .await?; info!( "clearing in-memory remote control enrollment because account id changed: websocket_url={}, previous_account_id={:?}, current_account_id={:?}", remote_control_target.websocket_url, enrollment .as_ref() .map(|enrollment| enrollment.account_id.as_str()), auth.account_id ); *enrollment = None; status_publisher.publish_environment_id(/*environment_id*/ None); if !connect_options.desired_state_tx.borrow().is_enabled() { return Err(io::Error::new( ErrorKind::Interrupted, "remote control disabled after account changed", )); } } if let Some(enrollment) = enrollment.as_mut() { enrollment.remote_control_target = remote_control_target.clone(); } if let Some(enrollment) = enrollment.as_ref() { status_publisher.publish_environment_id(Some(enrollment.environment_id.clone())); } if enrollment.is_none() { let loaded_enrollment = load_persisted_remote_control_enrollment( Some(state_db), remote_control_target, &auth.account_id, connect_options.app_server_client_name, ) .await?; if let Some(loaded_enrollment) = loaded_enrollment.as_ref() { status_publisher.publish_environment_id(Some(loaded_enrollment.environment_id.clone())); } *enrollment = loaded_enrollment.map(|mut enrollment| { enrollment.server_name = connect_options.server_name.to_string(); enrollment }); } enroll_and_persist_remote_control_server( remote_control_target, state_db, RemoteControlEnrollmentAuthContext { auth: &auth, recovery: auth_context, }, enrollment, connect_options, status_publisher, RemoteControlEnrollmentSelection::ReuseOrCreate, ) .await?; if enrollment .as_ref() .ok_or_else(|| io::Error::other("missing remote control enrollment after enrollment step"))? .should_refresh_server_token() { let enrollment_ref = enrollment.as_ref().ok_or_else(|| { io::Error::other("missing remote control enrollment after enrollment step") })?; let server_id = enrollment_ref.server_id.clone(); let environment_id = enrollment_ref.environment_id.clone(); info!( "refreshing remote control server token: websocket_url={}, refresh_url={}, account_id={}, server_id={}, environment_id={}", remote_control_target.websocket_url, remote_control_target.refresh_url, auth.account_id, server_id, environment_id ); let enrollment_ref = enrollment.as_mut().ok_or_else(|| { io::Error::other("missing remote control enrollment before server refresh") })?; match refresh_remote_control_server(&auth, connect_options.installation_id, enrollment_ref) .await { Ok(()) => {} Err(err) if err.kind() == ErrorKind::NotFound => { info!( "remote control server refresh returned HTTP 404; replacing stale enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}", remote_control_target.websocket_url, auth.account_id, server_id, environment_id ); enroll_and_persist_remote_control_server( remote_control_target, state_db, RemoteControlEnrollmentAuthContext { auth: &auth, recovery: auth_context, }, enrollment, connect_options, status_publisher, RemoteControlEnrollmentSelection::ReplaceExisting, ) .await?; } Err(err) if err.kind() == ErrorKind::PermissionDenied && recover_remote_control_auth( auth_context.auth_recovery, auth_context.auth_change_rx, ) .await => { return Err(io::Error::other(format!( "{err}; retrying after auth recovery" ))); } Err(err) => return Err(err), } } Ok(auth)}async fn resolve_desired_state_after_account_change( state_db: &StateRuntime, remote_control_target: &RemoteControlTarget, auth_manager: &Arc<AuthManager>, account_id: &str, connect_options: RemoteControlConnectOptions<'_>,) -> io::Result<()> { let durable_enabled = RemoteControlDesiredState::Enabled { persistence_preference: Some(true), }; if *connect_options.desired_state_tx.borrow() != durable_enabled { return Ok(()); } let _persistence = acquire_persistence_lock(connect_options.desired_state_persistence_lock).await; if *connect_options.desired_state_tx.borrow() != durable_enabled { return Ok(()); } let enrollment = state_db .get_remote_control_enrollment( &remote_control_target.websocket_url, account_id, connect_options.app_server_client_name, ) .await .map_err(io::Error::other)?; let current_auth = load_remote_control_auth(auth_manager).await?; if current_auth.account_id != account_id { return Err(io::Error::new( ErrorKind::WouldBlock, "remote control account changed while resolving persisted preference", )); } let resolved_state = desired_state_from_persisted_enrollment(enrollment); connect_options.desired_state_tx.send_if_modified(|state| { if *state != durable_enabled || *state == resolved_state { return false; } *state = resolved_state; true }); Ok(())}fn websocket_response_reports_missing_remote_app_server( response: &tungstenite::http::Response<Option<Vec<u8>>>,) -> bool { response.status().as_u16() == 404 && response.body().as_deref().is_some_and(|body| { serde_json::from_slice::<serde_json::Value>(body).is_ok_and(|body| { body.get("detail").and_then(serde_json::Value::as_str) == Some(REMOTE_APP_SERVER_NOT_FOUND_DETAIL) }) })}async fn replace_remote_control_enrollment_if_matches( state_db: Option<&StateRuntime>, remote_control_target: &RemoteControlTarget, auth_context: RemoteControlEnrollmentAuthContext<'_, '_>, current_enrollment: &CurrentRemoteControlEnrollment, enrollment: &RemoteControlEnrollment, connect_options: RemoteControlConnectOptions<'_>, status_publisher: &RemoteControlStatusPublisher,) -> io::Result<()> { let Some(state_db) = state_db else { return Err(io::Error::new( ErrorKind::NotFound, "remote control requires sqlite state db", )); }; let mut current_enrollment = current_enrollment.lock().await; if !current_enrollment .as_ref() .is_some_and(|current| same_remote_control_enrollment(current, enrollment)) { return Ok(()); } enroll_and_persist_remote_control_server( remote_control_target, state_db, auth_context, &mut current_enrollment, connect_options, status_publisher, RemoteControlEnrollmentSelection::ReplaceExisting, ) .await}async fn clear_remote_control_server_token_if_matches( current_enrollment: &CurrentRemoteControlEnrollment, enrollment: &RemoteControlEnrollment,) -> io::Result<()> { let mut current_enrollment = current_enrollment.lock().await; current_enrollment .as_mut() .filter(|current| same_remote_control_enrollment(current, enrollment)) .ok_or_else(|| { io::Error::other("missing remote control enrollment after websocket auth failure") })? .clear_server_token(); Ok(())}async fn enroll_and_persist_remote_control_server( remote_control_target: &RemoteControlTarget, state_db: &StateRuntime, auth_context: RemoteControlEnrollmentAuthContext<'_, '_>, enrollment: &mut Option<RemoteControlEnrollment>, connect_options: RemoteControlConnectOptions<'_>, status_publisher: &RemoteControlStatusPublisher, selection: RemoteControlEnrollmentSelection,) -> io::Result<()> { match selection { RemoteControlEnrollmentSelection::ReuseOrCreate => { if enrollment.is_some() { return Ok(()); } } RemoteControlEnrollmentSelection::ReplaceExisting => {} } if !connect_options.desired_state_tx.borrow().is_enabled() { return Err(io::Error::new( ErrorKind::Interrupted, "remote control disabled before enrollment", )); } info!( "creating new remote control enrollment: websocket_url={}, enroll_url={}, account_id={}", remote_control_target.websocket_url, remote_control_target.enroll_url, auth_context.auth.account_id ); let new_enrollment = match enroll_remote_control_server( remote_control_target, auth_context.auth, connect_options.installation_id, connect_options.server_name, ) .await { Ok(new_enrollment) => new_enrollment, Err(err) if err.kind() == ErrorKind::PermissionDenied && recover_remote_control_auth( auth_context.recovery.auth_recovery, auth_context.recovery.auth_change_rx, ) .await => { return Err(io::Error::other(format!( "{err}; retrying after auth recovery" ))); } Err(err) => return Err(err), }; let _persistence = acquire_persistence_lock(connect_options.desired_state_persistence_lock).await; let persistence_preference = match *connect_options.desired_state_tx.borrow() { RemoteControlDesiredState::Enabled { persistence_preference, } => persistence_preference, RemoteControlDesiredState::Unknown | RemoteControlDesiredState::Disabled => { return Err(io::Error::new( ErrorKind::Interrupted, "remote control disabled during enrollment", )); } }; if let Err(err) = update_persisted_remote_control_enrollment( Some(state_db), remote_control_target, &auth_context.auth.account_id, connect_options.app_server_client_name, Some(&new_enrollment), persistence_preference, ) .await { return Err(io::Error::other(format!( "failed to persist remote control enrollment in sqlite state db: {err}" ))); } info!( "created new remote control enrollment: websocket_url={}, account_id={}, server_id={}, environment_id={}", remote_control_target.websocket_url, new_enrollment.account_id, new_enrollment.server_id, new_enrollment.environment_id ); status_publisher.publish_environment_id(Some(new_enrollment.environment_id.clone())); *enrollment = Some(new_enrollment); Ok(())}fn format_remote_control_websocket_connect_error( websocket_url: &str, err: &tungstenite::Error,) -> String { let mut message = format!("failed to connect app-server remote control websocket `{websocket_url}`: {err}"); let tungstenite::Error::Http(response) = err else { return message; }; message.push_str(&format!(", {}", format_headers(response.headers()))); if let Some(body) = response.body().as_ref() && !body.is_empty() { let body_preview = preview_remote_control_response_body(body); message.push_str(&format!(", body: {body_preview}")); } message}#[cfg(test)]mod tests { use super::*; use crate::outgoing_message::OutgoingMessage; use crate::transport::remote_control::ServerEvent; use crate::transport::remote_control::auth::mark_recovery_auth_change_seen; use crate::transport::remote_control::protocol::StreamId; use crate::transport::remote_control::protocol::normalize_remote_control_url; use chrono::Utc; use codex_app_server_protocol::AuthMode; use codex_app_server_protocol::ConfigWarningNotification; use codex_app_server_protocol::JSONRPCMessage; use codex_app_server_protocol::JSONRPCNotification; use codex_app_server_protocol::ServerNotification; use codex_config::types::AuthCredentialsStoreMode; use codex_core::test_support::auth_manager_from_auth; use codex_login::AuthDotJson; use codex_login::AuthKeyringBackendKind; use codex_login::CodexAuth; use codex_login::save_auth; use codex_login::token_data::TokenData; use codex_login::token_data::parse_chatgpt_jwt_claims; use codex_state::StateRuntime; use futures::StreamExt; use pretty_assertions::assert_eq; use std::sync::Arc; use tempfile::TempDir; use tokio::io::AsyncBufReadExt; use tokio::io::AsyncWriteExt; use tokio::io::BufReader; use tokio::net::TcpListener; use tokio::net::TcpStream; use tokio::sync::mpsc; use tokio::time::Duration; use tokio::time::timeout; use tokio_tungstenite::accept_async; // Windows Bazel CI can take longer than a few seconds for the websocket // client connection attempt to reach the local test listener. #[cfg(windows)] const TEST_HTTP_ACCEPT_TIMEOUT: Duration = Duration::from_secs(30); #[cfg(not(windows))] const TEST_HTTP_ACCEPT_TIMEOUT: Duration = Duration::from_secs(5); const TEST_INSTALLATION_ID: &str = "11111111-1111-4111-8111-111111111111"; const TEST_REMOTE_CONTROL_SERVER_TOKEN: &str = "Remote Control Token"; fn remote_control_enrollment(remote_control_token: Option<&str>) -> RemoteControlEnrollment { RemoteControlEnrollment { remote_control_target: normalize_remote_control_url("http://localhost/backend-api/") .expect("target should normalize"), account_id: "account_id".to_string(), environment_id: "env_test".to_string(), server_id: "srv_e_test".to_string(), server_name: "test-server".to_string(), remote_control_token: remote_control_token.map(str::to_string), expires_at: remote_control_token .map(|_| time::OffsetDateTime::now_utc() + time::Duration::hours(1)), } } fn test_current_enrollment( enrollment: Option<RemoteControlEnrollment>, ) -> CurrentRemoteControlEnrollment { Arc::new(RemoteControlEnrollmentState::new(enrollment)) } #[test] fn next_reconnect_delay_resets_after_cap() { let mut reconnect_attempt = 9; let (reconnect_delay, reconnect_backoff_reset) = next_reconnect_delay(&mut reconnect_attempt); assert_eq!(reconnect_delay, REMOTE_CONTROL_RECONNECT_BACKOFF_CAP); assert!(reconnect_backoff_reset); assert_eq!(reconnect_attempt, 0); let (reconnect_delay, reconnect_backoff_reset) = next_reconnect_delay(&mut reconnect_attempt); assert!(reconnect_delay >= Duration::from_millis(180)); assert!(reconnect_delay <= Duration::from_millis(220)); assert!(!reconnect_backoff_reset); assert_eq!(reconnect_attempt, 1); } #[test] fn websocket_404_only_reports_explicit_missing_remote_app_server() { let cases = [ ( Some(br#"{"detail":"Remote app server not found"}"#.to_vec()), true, ), ( Some(br#" { "detail": "Remote app server not found", "extra": true } "#.to_vec()), true, ), (Some(br#"{"detail":"Not Found"}"#.to_vec()), false), (Some(b"Not Found".to_vec()), false), (Some(b"{".to_vec()), false), (Some(Vec::new()), false), (None, false), ]; for (body, expected) in cases { let response = tungstenite::http::Response::builder() .status(/*status*/ 404) .body(body) .expect("response should build"); assert_eq!( websocket_response_reports_missing_remote_app_server(&response), expected ); } let response = tungstenite::http::Response::builder() .status(/*status*/ 503) .body(Some( br#"{"detail":"Remote app server not found"}"#.to_vec(), )) .expect("response should build"); assert!(!websocket_response_reports_missing_remote_app_server( &response )); } fn remote_control_status_channel() -> ( RemoteControlStatusPublisher, watch::Receiver<RemoteControlStatusChangedNotification>, ) { let (status_tx, status_rx) = watch::channel(RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connecting, server_name: "test-server".to_string(), installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: None, }); (RemoteControlStatusPublisher::new(status_tx), status_rx) } fn enabled_desired_state_sender() -> watch::Sender<RemoteControlDesiredState> { watch::channel(RemoteControlDesiredState::Enabled { persistence_preference: None, }) .0 } #[test] fn mark_recovery_auth_change_seen_marks_only_recovery_revision_seen() { let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64); let auth_change_revision_before_recovery = *auth_change_rx.borrow(); auth_change_tx.send_modify(|revision| *revision += 1); mark_recovery_auth_change_seen(&mut auth_change_rx, auth_change_revision_before_recovery); assert!( !auth_change_rx .has_changed() .expect("auth change watch should remain open") ); } #[test] fn mark_recovery_auth_change_seen_preserves_racing_auth_change() { let (auth_change_tx, mut auth_change_rx) = watch::channel(0u64); let auth_change_revision_before_recovery = *auth_change_rx.borrow(); auth_change_tx.send_modify(|revision| *revision += 1); auth_change_tx.send_modify(|revision| *revision += 1); mark_recovery_auth_change_seen(&mut auth_change_rx, auth_change_revision_before_recovery); assert!( auth_change_rx .has_changed() .expect("auth change watch should remain open") ); } async fn remote_control_state_runtime(codex_home: &TempDir) -> Arc<StateRuntime> { StateRuntime::init(codex_home.path().to_path_buf(), "test-provider".to_string()) .await .expect("state runtime should initialize") } fn remote_control_auth_manager() -> Arc<AuthManager> { auth_manager_from_auth(CodexAuth::create_dummy_chatgpt_auth_for_testing()) } fn remote_control_url_for_listener(listener: &TcpListener) -> String { let addr = listener .local_addr() .expect("listener should have a local addr"); format!("http://{addr}/backend-api/") } fn remote_control_auth_dot_json(access_token: &str) -> AuthDotJson { #[derive(serde::Serialize)] struct Header { alg: &'static str, typ: &'static str, } let header = Header { alg: "none", typ: "JWT", }; let payload = serde_json::json!({ "email": "user@example.com", "https://api.openai.com/auth": { "chatgpt_user_id": "user-12345", "user_id": "user-12345", "chatgpt_account_id": "account_id" } }); let b64 = |bytes: &[u8]| base64::engine::general_purpose::URL_SAFE_NO_PAD.encode(bytes); let header_b64 = b64(&serde_json::to_vec(&header).expect("header should serialize")); let payload_b64 = b64(&serde_json::to_vec(&payload).expect("payload should serialize")); let fake_jwt = format!("{header_b64}.{payload_b64}.sig"); AuthDotJson { auth_mode: Some(AuthMode::Chatgpt), openai_api_key: None, tokens: Some(TokenData { id_token: parse_chatgpt_jwt_claims(&fake_jwt).expect("fake jwt should parse"), access_token: access_token.to_string(), refresh_token: "refresh-token".to_string(), account_id: Some("account_id".to_string()), }), last_refresh: Some(Utc::now()), agent_identity: None, personal_access_token: None, bedrock_api_key: None, } } #[tokio::test] async fn connect_remote_control_websocket_includes_http_error_details() { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let remote_control_url = remote_control_url_for_listener(&listener); let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); let expected_error = format!( "failed to connect app-server remote control websocket `{}`: HTTP error: 503 Service Unavailable, request-id: <none>, cf-ray: <none>, body: upstream unavailable", remote_control_target.websocket_url ); let server_task = tokio::spawn(async move { let (stream, request_line) = accept_http_request(&listener).await; assert_eq!( request_line, "GET /backend-api/wham/remote/control/server HTTP/1.1" ); respond_with_status_and_headers( stream, "503 Service Unavailable", &[("x-trace-id", "trace-503"), ("x-region", "us-east-1")], "upstream unavailable", ) .await; }); let codex_home = TempDir::new().expect("temp dir should create"); let state_db = remote_control_state_runtime(&codex_home).await; let auth_manager = remote_control_auth_manager(); let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some( TEST_REMOTE_CONTROL_SERVER_TOKEN, )))); let (status_publisher, status_rx) = remote_control_status_channel(); let err = match connect_remote_control_websocket( &remote_control_target, Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, desired_state_tx: &enabled_desired_state_sender(), desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) .await { Ok(_) => panic!("http error response should fail the websocket connect"), Err(err) => err, }; server_task.await.expect("server task should succeed"); assert_eq!(err.to_string(), expected_error); assert!(current_enrollment.lock().await.is_some()); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connecting, server_name: "test-server".to_string(), installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: Some("env_test".to_string()), } ); } #[tokio::test] async fn connect_remote_control_websocket_invalidates_unauthorized_server_token() { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let remote_control_url = remote_control_url_for_listener(&listener); let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); let codex_home = TempDir::new().expect("temp dir should create"); let state_db = remote_control_state_runtime(&codex_home).await; let auth_manager = remote_control_auth_manager(); let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some( TEST_REMOTE_CONTROL_SERVER_TOKEN, )))); let (status_publisher, status_rx) = remote_control_status_channel(); let server_task = tokio::spawn(async move { let (stream, request_line) = accept_http_request(&listener).await; assert_eq!( request_line, "GET /backend-api/wham/remote/control/server HTTP/1.1" ); respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await; }); let err = connect_remote_control_websocket( &remote_control_target, Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, desired_state_tx: &enabled_desired_state_sender(), desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) .await .expect_err("unauthorized response should fail the websocket connect"); server_task.await.expect("server task should succeed"); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connecting, server_name: "test-server".to_string(), installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: Some("env_test".to_string()), } ); assert_eq!( err.to_string(), "remote control websocket auth failed with HTTP 401 Unauthorized; refreshing server token before reconnect" ); let mut expected_enrollment = remote_control_enrollment(/*remote_control_token*/ None); expected_enrollment.remote_control_target = remote_control_target; assert_eq!(*current_enrollment.lock().await, Some(expected_enrollment)); } #[tokio::test] async fn connect_remote_control_websocket_recovers_after_unauthorized_enrollment() { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let remote_control_url = remote_control_url_for_listener(&listener); let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); let enroll_url = remote_control_target.enroll_url.clone(); let server_task = tokio::spawn(async move { let (stream, request_line) = accept_http_request(&listener).await; assert_eq!( request_line, "POST /backend-api/wham/remote/control/server/enroll HTTP/1.1" ); respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await; }); let codex_home = TempDir::new().expect("temp dir should create"); save_auth( codex_home.path(), &remote_control_auth_dot_json("stale-token"), AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), ) .expect("stale auth should save"); let state_db = remote_control_state_runtime(&codex_home).await; let auth_manager = AuthManager::shared( codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); let current_enrollment = test_current_enrollment(/*enrollment*/ None); let (status_publisher, status_rx) = remote_control_status_channel(); save_auth( codex_home.path(), &remote_control_auth_dot_json("fresh-token"), AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), ) .expect("fresh auth should save"); let err = connect_remote_control_websocket( &remote_control_target, Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, desired_state_tx: &enabled_desired_state_sender(), desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) .await .expect_err("unauthorized enrollment should fail the websocket connect"); server_task.await.expect("server task should succeed"); assert!( !status_rx .has_changed() .expect("remote control status watch should remain open") ); assert_eq!( err.to_string(), format!( "remote control server enrollment failed at `{enroll_url}`: HTTP 401 Unauthorized, request-id: <none>, cf-ray: <none>, body: unauthorized; retrying after auth recovery" ) ); assert_eq!( auth_manager .auth() .await .expect("auth should remain available") .get_token() .expect("token should be readable"), "fresh-token" ); assert!( !auth_change_rx .has_changed() .expect("auth change watch should remain open"), "recovery's own auth reload should not wake the reconnect loop" ); } #[tokio::test] async fn connect_remote_control_websocket_recovers_after_unauthorized_refresh() { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let remote_control_url = remote_control_url_for_listener(&listener); let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); let refresh_url = remote_control_target.refresh_url.clone(); let server_task = tokio::spawn(async move { let (stream, request_line) = accept_http_request(&listener).await; assert_eq!( request_line, "POST /backend-api/wham/remote/control/server/refresh HTTP/1.1" ); respond_with_status_and_headers(stream, "401 Unauthorized", &[], "unauthorized").await; }); let codex_home = TempDir::new().expect("temp dir should create"); save_auth( codex_home.path(), &remote_control_auth_dot_json("stale-token"), AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), ) .expect("stale auth should save"); let state_db = remote_control_state_runtime(&codex_home).await; let auth_manager = AuthManager::shared( codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); let current_enrollment = test_current_enrollment(Some(remote_control_enrollment( /*remote_control_token*/ None, ))); let (status_publisher, status_rx) = remote_control_status_channel(); save_auth( codex_home.path(), &remote_control_auth_dot_json("fresh-token"), AuthCredentialsStoreMode::File, AuthKeyringBackendKind::default(), ) .expect("fresh auth should save"); let err = connect_remote_control_websocket( &remote_control_target, Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, desired_state_tx: &enabled_desired_state_sender(), desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) .await .expect_err("unauthorized refresh should fail the websocket connect"); server_task.await.expect("server task should succeed"); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connecting, server_name: "test-server".to_string(), installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: Some("env_test".to_string()), } ); assert_eq!( err.to_string(), format!( "remote control server refresh failed at `{refresh_url}`: HTTP 401 Unauthorized, request-id: <none>, cf-ray: <none>, body: unauthorized; retrying after auth recovery" ) ); assert_eq!( auth_manager .auth() .await .expect("auth should remain available") .get_token() .expect("token should be readable"), "fresh-token" ); assert!( !auth_change_rx .has_changed() .expect("auth change watch should remain open"), "recovery's own auth reload should not wake the reconnect loop" ); } #[tokio::test] async fn connect_remote_control_websocket_requires_sqlite_state_db() { let remote_control_target = normalize_remote_control_url("http://127.0.0.1:9/backend-api/") .expect("target should parse"); let auth_manager = remote_control_auth_manager(); let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some( TEST_REMOTE_CONTROL_SERVER_TOKEN, )))); let (status_publisher, _status_rx) = remote_control_status_channel(); let err = connect_remote_control_websocket( &remote_control_target, /*state_db*/ None, RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, desired_state_tx: &enabled_desired_state_sender(), desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) .await .expect_err("missing sqlite state db should fail remote control"); assert_eq!(err.kind(), ErrorKind::NotFound); assert_eq!(err.to_string(), "remote control requires sqlite state db"); assert_eq!(*current_enrollment.lock().await, None); } #[tokio::test] async fn connect_remote_control_websocket_requires_chatgpt_auth() { let remote_control_target = normalize_remote_control_url("http://127.0.0.1:9/backend-api/") .expect("target should parse"); let codex_home = TempDir::new().expect("temp dir should create"); let state_db = remote_control_state_runtime(&codex_home).await; let auth_manager = AuthManager::shared( codex_home.path().to_path_buf(), /*enable_codex_api_key_env*/ false, AuthCredentialsStoreMode::File, /*chatgpt_base_url*/ None, AuthKeyringBackendKind::default(), ) .await; let mut auth_recovery = auth_manager.unauthorized_recovery(); let mut auth_change_rx = auth_manager.auth_change_receiver(); let current_enrollment = test_current_enrollment(Some(remote_control_enrollment(Some( TEST_REMOTE_CONTROL_SERVER_TOKEN, )))); let (status_publisher, mut status_rx) = remote_control_status_channel(); status_publisher.publish_environment_id(Some("env_test".to_string())); status_rx .changed() .await .expect("remote control status watch should remain open"); let err = connect_remote_control_websocket( &remote_control_target, Some(state_db.as_ref()), RemoteControlAuthContext { auth_manager: &auth_manager, auth_recovery: &mut auth_recovery, auth_change_rx: &mut auth_change_rx, }, ¤t_enrollment, RemoteControlConnectOptions { installation_id: TEST_INSTALLATION_ID, server_name: "test-server", subscribe_cursor: None, app_server_client_name: None, desired_state_tx: &enabled_desired_state_sender(), desired_state_persistence_lock: &Semaphore::new(1), }, &status_publisher, ) .await .expect_err("missing auth should fail remote control"); assert_eq!(err.kind(), ErrorKind::PermissionDenied); assert_eq!( err.to_string(), "remote control requires ChatGPT authentication" ); assert_eq!(*current_enrollment.lock().await, None); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connecting, server_name: "test-server".to_string(), installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: None, } ); } #[tokio::test] async fn run_remote_control_websocket_loop_shutdown_cancels_reconnect_backoff() { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let remote_control_url = remote_control_url_for_listener(&listener); drop(listener); let remote_control_target = normalize_remote_control_url(&remote_control_url).expect("target should parse"); let (transport_event_tx, transport_event_rx) = mpsc::channel(1); drop(transport_event_rx); let (status_publisher, _status_rx) = remote_control_status_channel(); let shutdown_token = CancellationToken::new(); let (desired_state_tx, _desired_state_rx) = watch::channel(RemoteControlDesiredState::Enabled { persistence_preference: None, }); let websocket_task = tokio::spawn({ let shutdown_token = shutdown_token.clone(); async move { RemoteControlWebsocket::new( RemoteControlWebsocketConfig { remote_control_url, installation_id: TEST_INSTALLATION_ID.to_string(), remote_control_target: Some(remote_control_target), server_name: "test-server".to_string(), }, /*state_db*/ None, remote_control_auth_manager(), RemoteControlChannels { transport_event_tx, status_publisher, current_enrollment: test_current_enrollment(/*enrollment*/ None), pairing_persistence_key: watch::channel(None).0, desired_state_persistence_lock: Arc::new(Semaphore::new(1)), }, shutdown_token, Arc::new(desired_state_tx), ) .run(/*app_server_client_name_rx*/ None) .await } }); tokio::time::sleep(Duration::from_millis(50)).await; shutdown_token.cancel(); timeout(Duration::from_millis(100), websocket_task) .await .expect("shutdown should cancel reconnect backoff") .expect("websocket task should join"); } #[tokio::test] async fn publish_status_if_changed_sends_only_status_changes() { let (status_publisher, mut status_rx) = remote_control_status_channel(); status_publisher.publish_environment_id(/*environment_id*/ None); assert!( timeout(Duration::from_millis(20), status_rx.changed()) .await .is_err() ); status_publisher.publish_environment_id(Some("env_first".to_string())); status_rx .changed() .await .expect("remote control status watch should remain open"); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connecting, server_name: "test-server".to_string(), installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: Some("env_first".to_string()), } ); status_publisher.publish_environment_id(Some("env_first".to_string())); assert!( timeout(Duration::from_millis(20), status_rx.changed()) .await .is_err() ); status_publisher.publish_status(RemoteControlConnectionStatus::Connected); status_rx .changed() .await .expect("remote control status watch should remain open"); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connected, server_name: "test-server".to_string(), installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: Some("env_first".to_string()), } ); status_publisher.publish_environment_id(/*environment_id*/ None); status_rx .changed() .await .expect("remote control status watch should remain open"); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Connected, server_name: "test-server".to_string(), installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: None, } ); status_publisher.publish_environment_id(Some("env_disabled".to_string())); status_publisher.publish_status(RemoteControlConnectionStatus::Disabled); status_rx .changed() .await .expect("remote control status watch should remain open"); assert_eq!( status_rx.borrow().clone(), RemoteControlStatusChangedNotification { status: RemoteControlConnectionStatus::Disabled, server_name: "test-server".to_string(), installation_id: TEST_INSTALLATION_ID.to_string(), environment_id: None, } ); status_publisher.publish_environment_id(Some("env_disabled".to_string())); assert!( timeout(Duration::from_millis(20), status_rx.changed()) .await .is_err() ); } #[tokio::test] async fn run_server_writer_inner_sends_periodic_ping_frames() { let (client_stream, mut server_stream) = connected_websocket_pair().await; let (websocket_writer, _websocket_reader) = client_stream.split(); let (outbound_buffer, used_rx) = BoundedOutboundBuffer::new(); let state = Arc::new(Mutex::new(WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), })); let (_server_event_tx, server_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY); let server_event_rx = Arc::new(Mutex::new(server_event_rx)); let shutdown_token = CancellationToken::new(); let writer_task = tokio::spawn(RemoteControlWebsocket::run_server_writer_inner( state, server_event_rx, used_rx, websocket_writer, Duration::from_millis(20), shutdown_token.clone(), )); let message = timeout(Duration::from_secs(5), server_stream.next()) .await .expect("ping frame should arrive in time") .expect("server websocket should stay open") .expect("ping frame should read"); assert!(matches!(message, tungstenite::Message::Ping(_))); shutdown_token.cancel(); writer_task .await .expect("writer task should join") .expect("writer should stop cleanly"); } #[tokio::test] async fn join_connection_workers_aborts_stuck_worker_after_timeout() { let mut join_set = tokio::task::JoinSet::new(); join_set.spawn(futures::future::pending::<()>()); RemoteControlWebsocket::join_connection_workers(&mut join_set, Duration::from_millis(10)) .await; assert!(join_set.is_empty()); } #[tokio::test] async fn run_server_writer_inner_assigns_contiguous_seq_ids_per_stream() { let (client_stream, mut server_stream) = connected_websocket_pair().await; let (websocket_writer, _websocket_reader) = client_stream.split(); let (outbound_buffer, used_rx) = BoundedOutboundBuffer::new(); let state = Arc::new(Mutex::new(WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), })); let (server_event_tx, server_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY); let server_event_rx = Arc::new(Mutex::new(server_event_rx)); let shutdown_token = CancellationToken::new(); let writer_task = tokio::spawn(RemoteControlWebsocket::run_server_writer_inner( state, server_event_rx, used_rx, websocket_writer, Duration::from_secs(60), shutdown_token.clone(), )); let client_id = ClientId("client-1".to_string()); let first_stream = StreamId("stream-1".to_string()); let second_stream = StreamId("stream-2".to_string()); for stream_id in [&first_stream, &second_stream, &first_stream] { server_event_tx .send(super::super::QueuedServerEnvelope { event: ServerEvent::Pong { status: crate::transport::remote_control::protocol::PongStatus::Active, }, client_id: client_id.clone(), stream_id: stream_id.clone(), write_complete_tx: None, }) .await .expect("server event should queue"); } assert_eq!( read_server_text_event(&mut server_stream).await, serde_json::json!({ "type": "pong", "client_id": "client-1", "stream_id": "stream-1", "seq_id": 1, "status": "active", }) ); assert_eq!( read_server_text_event(&mut server_stream).await, serde_json::json!({ "type": "pong", "client_id": "client-1", "stream_id": "stream-2", "seq_id": 1, "status": "active", }) ); assert_eq!( read_server_text_event(&mut server_stream).await, serde_json::json!({ "type": "pong", "client_id": "client-1", "stream_id": "stream-1", "seq_id": 2, "status": "active", }) ); shutdown_token.cancel(); writer_task .await .expect("writer task should join") .expect("writer should stop cleanly"); } #[tokio::test] async fn run_websocket_reader_inner_times_out_without_pong_frames() { let (client_stream, _server_stream) = connected_websocket_pair().await; let (_websocket_writer, websocket_reader) = client_stream.split(); let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let state = Arc::new(Mutex::new(WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), })); let (server_event_tx, _server_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY); let (transport_event_tx, _transport_event_rx) = mpsc::channel(super::super::CHANNEL_CAPACITY); let shutdown_token = CancellationToken::new(); let client_tracker = Arc::new(Mutex::new(ClientTracker::new( server_event_tx, transport_event_tx, &shutdown_token, ))); let err = timeout( Duration::from_secs(5), RemoteControlWebsocket::run_websocket_reader_inner( client_tracker, state, websocket_reader, Duration::from_millis(100), shutdown_token, ), ) .await .expect("reader should time out waiting for pong") .expect_err("missing pong should fail the websocket reader"); assert_eq!(err.kind(), ErrorKind::TimedOut); assert_eq!(err.to_string(), "remote control websocket pong timeout"); } #[test] fn outbound_buffer_acks_by_stream_id() { let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new(); let client_1 = ClientId("client-1".to_string()); let client_2 = ClientId("client-2".to_string()); let stream_1 = StreamId("stream-1".to_string()); outbound_buffer.insert(&server_envelope( &client_1, "stream-1", /*seq_id*/ 1, "first-client-old-stream", )); outbound_buffer.insert(&server_envelope( &client_2, "stream-1", /*seq_id*/ 2, "second-client", )); outbound_buffer.insert(&server_envelope( &client_1, "stream-2", /*seq_id*/ 3, "first-client-new-stream", )); outbound_buffer.ack( &client_1, &stream_1, /*acked_seq_id*/ 3, /*acked_segment_id*/ None, ); let mut retained = outbound_buffer .server_envelopes() .map(|server_envelope| { ( server_envelope.client_id.0.as_str(), server_envelope.stream_id.0.as_str(), server_envelope.seq_id, ) }) .collect::<Vec<_>>(); retained.sort_unstable(); assert_eq!( retained, vec![("client-1", "stream-2", 3), ("client-2", "stream-1", 2)] ); assert_eq!(*used_rx.borrow(), 2); } #[test] fn outbound_buffer_retains_unacked_messages_until_ack_advances() { let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new(); let client_1 = ClientId("client-1".to_string()); let client_2 = ClientId("client-2".to_string()); let stream_1 = StreamId("stream-1".to_string()); outbound_buffer.insert(&server_envelope( &client_1, "stream-1", /*seq_id*/ 1, "first-old", )); outbound_buffer.insert(&server_envelope( &client_1, "stream-2", /*seq_id*/ 2, "first-new", )); outbound_buffer.insert(&server_envelope( &client_2, "stream-1", /*seq_id*/ 3, "second", )); outbound_buffer.ack( &client_1, &stream_1, /*acked_seq_id*/ 1, /*acked_segment_id*/ None, ); let mut retained = outbound_buffer .server_envelopes() .map(|server_envelope| { ( server_envelope.client_id.0.as_str(), server_envelope.stream_id.0.as_str(), server_envelope.seq_id, ) }) .collect::<Vec<_>>(); retained.sort_unstable(); assert_eq!( retained, vec![("client-1", "stream-2", 2), ("client-2", "stream-1", 3)] ); assert_eq!(*used_rx.borrow(), 2); } #[test] fn outbound_buffer_advances_segmented_acks_by_wire_cursor() { let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new(); let client_id = ClientId("client-1".to_string()); let stream_id = StreamId("stream-1".to_string()); outbound_buffer.insert(&server_chunk_envelope( &client_id, "stream-1", /*seq_id*/ 4, /*segment_id*/ 0, )); outbound_buffer.insert(&server_chunk_envelope( &client_id, "stream-1", /*seq_id*/ 4, /*segment_id*/ 1, )); outbound_buffer.ack( &client_id, &stream_id, /*acked_seq_id*/ 4, /*acked_segment_id*/ Some(1), ); let retained = outbound_buffer .server_envelopes() .map(|server_envelope| server_envelope.event.segment_id()) .collect::<Vec<_>>(); assert_eq!(retained, Vec::<Option<usize>>::new()); assert_eq!(*used_rx.borrow(), 0); } #[test] fn outbound_buffer_treats_segmentless_acks_as_seq_level_acks() { let (mut outbound_buffer, used_rx) = BoundedOutboundBuffer::new(); let client_id = ClientId("client-1".to_string()); let stream_id = StreamId("stream-1".to_string()); outbound_buffer.insert(&server_chunk_envelope( &client_id, "stream-1", /*seq_id*/ 4, /*segment_id*/ 0, )); outbound_buffer.insert(&server_chunk_envelope( &client_id, "stream-1", /*seq_id*/ 4, /*segment_id*/ 1, )); outbound_buffer.ack( &client_id, &stream_id, /*acked_seq_id*/ 4, /*acked_segment_id*/ None, ); let retained = outbound_buffer .server_envelopes() .map(|server_envelope| server_envelope.event.segment_id()) .collect::<Vec<_>>(); assert_eq!(retained, Vec::<Option<usize>>::new()); assert_eq!(*used_rx.borrow(), 0); } #[test] fn websocket_state_drops_duplicate_client_chunks_while_pending() { let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let mut state = WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), }; let first_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0, /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x", ); let second_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 1, /*segment_count*/ 2, /*message_size_bytes*/ 2, b"y", ); assert!(matches!( observe_client_message(&mut state, first_chunk.clone()), ClientSegmentObservation::Pending )); assert!(matches!( observe_client_message(&mut state, first_chunk.clone()), ClientSegmentObservation::Dropped )); assert!(matches!( observe_client_message(&mut state, second_chunk), ClientSegmentObservation::Dropped )); assert!(matches!( observe_client_message(&mut state, first_chunk), ClientSegmentObservation::Pending )); } #[test] fn websocket_state_drops_replayed_client_chunks_after_completion() { let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let mut state = WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), }; let message = JSONRPCMessage::Notification(JSONRPCNotification { method: "initialized".to_string(), params: None, }); let raw = serde_json::to_vec(&message).expect("message should serialize"); let split = raw.len() / 2; let first_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0, /*segment_count*/ 2, raw.len(), &raw[..split], ); let second_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 1, /*segment_count*/ 2, raw.len(), &raw[split..], ); assert!(matches!( observe_client_message(&mut state, first_chunk.clone()), ClientSegmentObservation::Pending )); let completed_envelope = match observe_client_message(&mut state, second_chunk) { ClientSegmentObservation::Forward(client_envelope) => *client_envelope, _ => panic!("expected completed client message"), }; state.record_client_message_delivery( &completed_envelope, Some(( ( ClientId("client-1".to_string()), Some(StreamId("stream-1".to_string())), ), 4, )), ); assert!(matches!( observe_client_message(&mut state, first_chunk), ClientSegmentObservation::Dropped )); } #[test] fn websocket_state_allows_replay_before_completed_chunk_delivery() { let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let mut state = WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), }; let message = JSONRPCMessage::Notification(JSONRPCNotification { method: "initialized".to_string(), params: None, }); let raw = serde_json::to_vec(&message).expect("message should serialize"); let split = raw.len() / 2; let first_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0, /*segment_count*/ 2, raw.len(), &raw[..split], ); let second_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 1, /*segment_count*/ 2, raw.len(), &raw[split..], ); assert!(matches!( observe_client_message(&mut state, first_chunk.clone()), ClientSegmentObservation::Pending )); assert!(matches!( observe_client_message(&mut state, second_chunk), ClientSegmentObservation::Forward(_) )); assert!(matches!( observe_client_message(&mut state, first_chunk), ClientSegmentObservation::Pending )); } #[test] fn websocket_state_allows_replay_after_rejected_out_of_order_chunk() { let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let mut state = WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), }; let first_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0, /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x", ); let second_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 1, /*segment_count*/ 2, /*message_size_bytes*/ 2, b"y", ); assert!(matches!( observe_client_message(&mut state, second_chunk), ClientSegmentObservation::Dropped )); assert!(matches!( observe_client_message(&mut state, first_chunk), ClientSegmentObservation::Pending )); } #[test] fn websocket_state_allows_replay_after_later_chunk_drops() { let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let mut state = WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), }; let first_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0, /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x", ); let invalid_second_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 1, /*segment_count*/ 2, /*message_size_bytes*/ 2, b"", ); assert!(matches!( observe_client_message(&mut state, first_chunk.clone()), ClientSegmentObservation::Pending )); assert!(matches!( observe_client_message(&mut state, invalid_second_chunk), ClientSegmentObservation::Dropped )); assert!(matches!( observe_client_message(&mut state, first_chunk), ClientSegmentObservation::Pending )); } #[test] fn websocket_state_drops_oversized_client_chunk_frames() { let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let mut state = WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), }; let chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0, /*segment_count*/ 1, /*message_size_bytes*/ 1, b"x", ); assert!(matches!( state.observe_client_message(chunk, REMOTE_CONTROL_SEGMENT_MAX_BYTES + 1), ClientSegmentObservation::Dropped )); } #[test] fn websocket_state_ignores_oversized_stale_chunks_without_dropping_newer_assembly() { let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let mut state = WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), }; let message = JSONRPCMessage::Notification(JSONRPCNotification { method: "initialized".to_string(), params: None, }); let raw = serde_json::to_vec(&message).expect("message should serialize"); let split = raw.len() / 2; let first_newer_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 8, /*segment_id*/ 0, /*segment_count*/ 2, raw.len(), &raw[..split], ); let oversized_stale_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 7, /*segment_id*/ 0, /*segment_count*/ 2, raw.len(), &raw[..split], ); let second_newer_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 8, /*segment_id*/ 1, /*segment_count*/ 2, raw.len(), &raw[split..], ); assert!(matches!( observe_client_message(&mut state, first_newer_chunk), ClientSegmentObservation::Pending )); assert!(matches!( state.observe_client_message( oversized_stale_chunk, REMOTE_CONTROL_SEGMENT_MAX_BYTES + 1, ), ClientSegmentObservation::Dropped )); assert!(matches!( observe_client_message(&mut state, second_newer_chunk), ClientSegmentObservation::Forward(_) )); } #[test] fn websocket_state_ignores_oversized_duplicate_chunks_without_dropping_current_assembly() { let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let mut state = WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), }; let message = JSONRPCMessage::Notification(JSONRPCNotification { method: "initialized".to_string(), params: None, }); let raw = serde_json::to_vec(&message).expect("message should serialize"); let split = raw.len() / 2; let first_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 8, /*segment_id*/ 0, /*segment_count*/ 2, raw.len(), &raw[..split], ); let oversized_duplicate_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 8, /*segment_id*/ 0, /*segment_count*/ 2, raw.len(), &raw[..split], ); let second_chunk = client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 8, /*segment_id*/ 1, /*segment_count*/ 2, raw.len(), &raw[split..], ); assert!(matches!( observe_client_message(&mut state, first_chunk), ClientSegmentObservation::Pending )); assert!(matches!( state.observe_client_message( oversized_duplicate_chunk, REMOTE_CONTROL_SEGMENT_MAX_BYTES + 1, ), ClientSegmentObservation::Dropped )); assert!(matches!( observe_client_message(&mut state, second_chunk), ClientSegmentObservation::Forward(_) )); } #[test] fn websocket_state_clears_chunk_cursor_when_stream_is_invalidated() { let (outbound_buffer, _used_rx) = BoundedOutboundBuffer::new(); let mut state = WebsocketState { outbound_buffer, subscribe_cursor: None, next_seq_id_by_stream: HashMap::new(), last_completed_client_chunk_seq_id_by_stream: HashMap::new(), client_segment_reassembler: ClientSegmentReassembler::default(), }; let client_id = ClientId("client-1".to_string()); let stream_id = StreamId("stream-1".to_string()); assert!(matches!( observe_client_message( &mut state, client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 4, /*segment_id*/ 0, /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x", ) ), ClientSegmentObservation::Pending )); state.invalidate_client_message_stream(&client_id, &stream_id); state .client_segment_reassembler .invalidate_stream(&client_id, &stream_id); assert!(matches!( observe_client_message( &mut state, client_chunk_envelope( "client-1", "stream-1", /*seq_id*/ 1, /*segment_id*/ 0, /*segment_count*/ 2, /*message_size_bytes*/ 2, b"x", ) ), ClientSegmentObservation::Pending )); } fn server_envelope( client_id: &ClientId, stream_id: &str, seq_id: u64, summary: &str, ) -> ServerEnvelope { ServerEnvelope { event: ServerEvent::ServerMessage { message: Box::new(OutgoingMessage::AppServerNotification( ServerNotification::ConfigWarning(ConfigWarningNotification { summary: summary.to_string(), details: None, path: None, range: None, }), )), }, client_id: client_id.clone(), stream_id: StreamId(stream_id.to_string()), seq_id, } } fn server_chunk_envelope( client_id: &ClientId, stream_id: &str, seq_id: u64, segment_id: usize, ) -> ServerEnvelope { ServerEnvelope { event: ServerEvent::ServerMessageChunk { segment_id, segment_count: 2, message_size_bytes: 2, message_chunk_base64: String::new(), }, client_id: client_id.clone(), stream_id: StreamId(stream_id.to_string()), seq_id, } } fn client_chunk_envelope( client_id: &str, stream_id: &str, seq_id: u64, segment_id: usize, segment_count: usize, message_size_bytes: usize, chunk: &[u8], ) -> ClientEnvelope { ClientEnvelope { event: ClientEvent::ClientMessageChunk { segment_id, segment_count, message_size_bytes, message_chunk_base64: base64::engine::general_purpose::STANDARD.encode(chunk), }, client_id: ClientId(client_id.to_string()), stream_id: Some(StreamId(stream_id.to_string())), seq_id: Some(seq_id), cursor: None, } } fn observe_client_message( state: &mut WebsocketState, envelope: ClientEnvelope, ) -> ClientSegmentObservation { let wire_size_bytes = serde_json::to_vec(&envelope) .expect("client envelope should serialize") .len(); state.observe_client_message(envelope, wire_size_bytes) } async fn accept_http_request(listener: &TcpListener) -> (TcpStream, String) { let (stream, _) = timeout(TEST_HTTP_ACCEPT_TIMEOUT, listener.accept()) .await .expect("HTTP request should arrive in time") .expect("listener accept should succeed"); let mut reader = BufReader::new(stream); let mut request_line = String::new(); reader .read_line(&mut request_line) .await .expect("request line should read"); loop { let mut line = String::new(); reader .read_line(&mut line) .await .expect("header line should read"); if line == "\r\n" { break; } } ( reader.into_inner(), request_line.trim_end_matches("\r\n").to_string(), ) } async fn connected_websocket_pair() -> ( WebSocketStream<MaybeTlsStream<TcpStream>>, WebSocketStream<TcpStream>, ) { let listener = TcpListener::bind("127.0.0.1:0") .await .expect("listener should bind"); let connect_task = tokio::spawn(connect_async(format!( "ws://{}", listener .local_addr() .expect("listener should have a local addr") ))); let (server_stream, _) = listener .accept() .await .expect("server should accept client"); let server_stream = accept_async(server_stream) .await .expect("server websocket handshake should succeed"); let (client_stream, _) = connect_task .await .expect("client connect task should join") .expect("client websocket handshake should succeed"); (client_stream, server_stream) } async fn read_server_text_event( server_stream: &mut WebSocketStream<TcpStream>, ) -> serde_json::Value { let message = timeout(Duration::from_secs(5), server_stream.next()) .await .expect("server event should arrive in time") .expect("server websocket should stay open") .expect("server event should read"); let tungstenite::Message::Text(text) = message else { panic!("expected text event, got {message:?}"); }; serde_json::from_str(text.as_ref()).expect("server event should deserialize") } async fn respond_with_status_and_headers( mut stream: TcpStream, status: &str, headers: &[(&str, &str)], body: &str, ) { let extra_headers = headers .iter() .map(|(name, value)| format!("{name}: {value}\r\n")) .collect::<String>(); let response = format!( "HTTP/1.1 {status}\r\ncontent-type: text/plain\r\ncontent-length: {}\r\nconnection: close\r\n{extra_headers}\r\n{body}", body.len(), ); stream .write_all(response.as_bytes()) .await .expect("response should write"); stream.flush().await.expect("response should flush"); }}