Skip to main content

nautilus_network/websocket/
client.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! Connection lifecycle and task coordination for [`WebSocketClient`].
17//!
18//! # Execution model
19//!
20//! Handler mode runs a controller, reader, and writer task, plus an optional heartbeat task. The
21//! reader dispatches frames, the writer serializes all sink access, and the controller drives
22//! reconnect and shutdown from the shared [`ConnectionMode`]. Stream mode omits the managed
23//! reader task and automatic reconnect.
24//!
25//! A read failure ends the reader and wakes the controller; a write failure requests reconnect and
26//! also wakes it. The controller establishes a replacement transport with exponential backoff,
27//! hands its sink to the writer, replaces the reader, then publishes the reconnect notification.
28//! Terminal disconnect and close states stop all tasks instead of starting another reconnect.
29//!
30//! # Reconnect ordering
31//!
32//! An accepted reconnect request first transitions to `Reconnect` and invalidates registered
33//! authentication. The controller may establish the replacement while the synchronous loss
34//! callback runs, but cannot restore `Active` until loss publication completes. A reconnect through
35//! [`WebSocketReconnectHandle`] releases its controller request fence before invoking the callback,
36//! so dropping the client can terminate controller work even if the callback blocks. Terminal
37//! states always take precedence.
38//!
39//! # Send semantics
40//!
41//! Ordinary sends return after entering the writer channel. Failed ordinary writes and ordinary
42//! messages encountered during reconnect enter a FIFO buffer for replay on an active replacement
43//! connection. While that replay is pending, ordinary sends also enter the buffer and queue behind
44//! it instead of reaching the replacement sink early. Registered authentication state can hold or
45//! discard that replay before it reaches the sink. Control frames are never buffered, so a pong
46//! sent while replay waits still writes to the replacement sink directly.
47//! [`WebSocketConfig::writer_capacity`] bounds ordinary messages across the channel, in-flight
48//! writes, and replay. Full capacity rejects a send with [`SendError::BufferFull`] before
49//! enqueueing. Ownership-bound sends and control traffic use a separate bounded allowance so
50//! authentication can unblock replay.
51//!
52//! Ownership-bound sends wait for the writer result, require the expected connection epoch, and
53//! never enter the reconnect buffer. The writer advances the epoch only when it installs a
54//! replacement sink, so inbound attribution and bound sends share one transport ownership boundary.
55
56use std::{
57    collections::VecDeque,
58    fmt::Debug,
59    future::Future,
60    pin::pin,
61    sync::{
62        Arc, OnceLock,
63        atomic::{AtomicBool, AtomicU8, AtomicU32, AtomicU64, Ordering},
64    },
65    time::Duration,
66};
67
68use futures_util::{SinkExt, StreamExt};
69use http::HeaderName;
70use nautilus_core::string::secret::REDACTED;
71use nautilus_cryptography::providers::install_cryptographic_provider;
72use parking_lot::RwLock;
73#[cfg(any(feature = "turmoil", feature = "transport-sockudo"))]
74use rustls::ClientConfig;
75#[cfg(feature = "transport-sockudo")]
76use sockudo_ws::{
77    Config as SockudoConfig, Http1, Role, Stream as SockudoStream,
78    WebSocketStream as SockudoWebSocketStream,
79};
80#[cfg(feature = "transport-sockudo")]
81use tokio::io::{AsyncRead, AsyncWrite};
82#[cfg(any(feature = "turmoil", feature = "transport-sockudo"))]
83use tokio_rustls::TlsConnector;
84#[cfg(feature = "turmoil")]
85use tokio_tungstenite::MaybeTlsStream;
86#[cfg(feature = "turmoil")]
87use tokio_tungstenite::client_async_with_config;
88#[cfg(all(not(feature = "turmoil"), not(all(feature = "simulation", madsim))))]
89use tokio_tungstenite::connect_async_with_config;
90use tokio_tungstenite::tungstenite::{
91    client::IntoClientRequest, handshake::client::Request, http::HeaderValue,
92};
93use tokio_util::sync::CancellationToken;
94use ustr::Ustr;
95
96#[cfg(not(feature = "turmoil"))]
97use super::proxy::{ProxyKind, WsTarget, tunnel_via_proxy};
98use super::{
99    auth::{AuthState, AuthTracker},
100    config::{InitialConnectRetryPolicy, TransportBackend, WebSocketConfig},
101    consts::{
102        CONNECTION_STATE_CHECK_INTERVAL_MS, GRACEFUL_SHUTDOWN_DELAY_MS,
103        GRACEFUL_SHUTDOWN_TIMEOUT_SECS,
104    },
105    types::{
106        CloseHandler, EpochMessageHandler, EpochPingHandler, MessageHandler, MessageReader,
107        MessageWriter, PingHandler, WriterCommand,
108    },
109};
110#[cfg(feature = "turmoil")]
111use crate::net::TcpConnector;
112#[cfg(feature = "transport-sockudo")]
113use crate::net::TcpStream;
114#[cfg(feature = "transport-sockudo")]
115use crate::transport::sockudo::{
116    PrefixedIo, SockudoTransport, client_handshake_with_headers, validate_extra_headers,
117};
118use crate::{
119    RECONNECTED, SocketState, SocketStateSink,
120    backoff::{
121        ExponentialBackoff, RECONNECT_STABILITY_THRESHOLD, ReconnectThrottle, wait_reconnect_delay,
122    },
123    dst,
124    error::{SendError, is_connection_drop_io_error},
125    logging::{log_task_aborted, log_task_started, log_task_stopped},
126    mode::{
127        ConnectionMode, ControllerLifecycle, ReadSessionFence, ReconnectOutcome,
128        ReconnectRequestOutcome,
129    },
130    ratelimiter::{RateLimiter, clock::MonotonicClock, quota::Quota},
131    retry::{RetryConfig, RetryError, RetryManager},
132    transport::{
133        BoxedWsTransport, CloseFrame, Message, TransportError, error::retryable_status,
134        tungstenite::TungsteniteTransport,
135    },
136    writer::{DEFAULT_WRITER_CAPACITY, WriterReceiver, WriterSender},
137};
138
139const WRITE_TIMEOUT_SECS: u64 = 5;
140const CONTROLLER_FALLBACK_INTERVAL_MS: u64 = 100;
141const INITIAL_CONNECT_OPERATION: &str = "WebSocket initial connection";
142
143/// The RFC 6455 control-frame payload limit.
144const MAX_CONTROL_FRAME_PAYLOAD_BYTES: usize = 125;
145
146/// Owns the transport tasks and reconnect state used by [`WebSocketClient`].
147///
148/// # Connection ownership
149///
150/// The client uses one reader and supports concurrent senders. In handler mode, a reader task
151/// dispatches incoming messages while a writer task serializes sends received over a channel. The
152/// controller owns the connection lifecycle and replaces both transport halves during reconnects.
153///
154/// Stream mode returns the reader to the caller. The client cannot replace that reader, so stream
155/// mode disables automatic reconnection.
156///
157/// # Heartbeats
158///
159/// When configured, a dedicated task sends heartbeat messages at the requested interval. Configure
160/// the interval below the server's heartbeat deadline. Handler mode can also opt into a heartbeat
161/// timeout that reconnects when no frame arrives within a venue-specific duration. The timeout
162/// starts with each connection and resets on every inbound frame, including Ping and Pong.
163///
164/// # Reconnection
165///
166/// The writer task owns queued sends across reconnects. A successful reconnect installs the
167/// replacement writer and starts a reader for the new connection epoch. Depending on the
168/// configured authentication gate, buffered sends drain immediately or wait for the new session to
169/// authenticate. Failed authentication discards messages that remain buffered.
170pub struct WebSocketClientInner {
171    config: WebSocketConfig,
172    reconnect_headers: ReconnectHeaders,
173    handler: Option<IncomingHandler>,
174    ping_handler: Option<IncomingPingHandler>,
175    close_handler: Option<CloseHandler>,
176    read_task: Option<dst::task::JoinHandle<()>>,
177    read_fence: Option<ReadSessionFence>,
178    write_task: dst::task::JoinHandle<()>,
179    writer_tx: WriterSender<WriterCommand>,
180    heartbeat_task: Option<dst::task::JoinHandle<()>>,
181    connection_mode: Arc<AtomicU8>,
182    connection_epoch: Arc<AtomicU64>,
183    state_notify: Arc<tokio::sync::Notify>,
184    controller_notify: Arc<tokio::sync::Notify>,
185    reconnect_published: Arc<AtomicBool>,
186    connect_timeout: Duration,
187    heartbeat_timeout: Option<Duration>,
188    backoff: ExponentialBackoff,
189    reconnect_throttle: ReconnectThrottle,
190    reconnect_max_attempts: Option<u32>,
191    reconnection_attempt_count: u32,
192    auth_tracker: Arc<OnceLock<AuthTracker>>,
193    reconnect_buffer_waits_for_auth: Arc<AtomicBool>,
194    state_sink: Option<SocketStateSink>,
195    connection_rate_limit: Option<ConnectionRateLimit>,
196}
197
198#[derive(Clone, Debug)]
199struct ConnectionRateLimit {
200    limiter: Arc<RateLimiter<Ustr, MonotonicClock>>,
201    keys: Arc<[Ustr]>,
202}
203
204#[derive(Default)]
205struct InitialConnectOptions {
206    retry_policy: Option<InitialConnectRetryPolicy>,
207    cancellation_token: Option<CancellationToken>,
208}
209
210impl WebSocketClientInner {
211    /// Creates an inner WebSocket client with an existing writer.
212    ///
213    /// This is used for stream mode where the reader is owned by the caller.
214    ///
215    /// # Errors
216    ///
217    /// Returns an error if the exponential backoff configuration is invalid.
218    #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
219    #[expect(
220        clippy::unused_async,
221        clippy::unused_async_trait_impl,
222        reason = "async signature for consistency with connect-based constructors"
223    )]
224    pub async fn new_with_writer(
225        config: WebSocketConfig,
226        writer: MessageWriter,
227    ) -> Result<Self, TransportError> {
228        Self::new_with_writer_and_state_sink(config, writer, None)
229    }
230
231    fn new_with_writer_and_state_sink(
232        mut config: WebSocketConfig,
233        writer: MessageWriter,
234        state_sink: Option<SocketStateSink>,
235    ) -> Result<Self, TransportError> {
236        install_cryptographic_provider();
237
238        crate::writer::validate_capacity(config.writer_capacity).map_err(|e| {
239            TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
240        })?;
241
242        if config.heartbeat_interval_secs == Some(0) {
243            return Err(TransportError::Io(std::io::Error::new(
244                std::io::ErrorKind::InvalidInput,
245                "Heartbeat interval cannot be zero",
246            )));
247        }
248
249        let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
250        let connection_epoch = Arc::new(AtomicU64::new(0));
251        let state_notify = Arc::new(tokio::sync::Notify::new());
252        let controller_notify = Arc::new(tokio::sync::Notify::new());
253        let reconnect_published = Arc::new(AtomicBool::new(true));
254        let outcome =
255            ConnectionMode::complete_reconnect_with_sink(&connection_mode, state_sink.as_ref());
256        debug_assert_eq!(outcome, ReconnectOutcome::Reconnected);
257
258        // Note: We don't spawn a read task here since the reader is handled externally
259        let read_task = None;
260        let read_fence = None;
261
262        // Stream mode ignores reconnect settings, use harmless defaults
263        let backoff = ExponentialBackoff::new(
264            Duration::from_secs(2),
265            Duration::from_secs(30),
266            1.5,
267            100,
268            true,
269        )
270        .map_err(|e| {
271            TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
272        })?;
273
274        let auth_tracker = Arc::new(OnceLock::new());
275        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
276
277        let (writer_tx, writer_rx) = crate::writer::channel::<WriterCommand>(
278            config.writer_capacity.unwrap_or(DEFAULT_WRITER_CAPACITY),
279        );
280        let write_task = Self::spawn_write_task(
281            Arc::clone(&connection_mode),
282            Arc::clone(&controller_notify),
283            Arc::clone(&reconnect_published),
284            writer,
285            writer_rx,
286            Arc::clone(&connection_epoch),
287            Arc::clone(&auth_tracker),
288            Arc::clone(&reconnect_buffer_waits_for_auth),
289            state_sink.clone(),
290        );
291
292        let heartbeat_task = config.heartbeat_interval_secs.map(|heartbeat_interval| {
293            Self::spawn_heartbeat_task(
294                Arc::clone(&connection_mode),
295                heartbeat_interval,
296                config.heartbeat_payload.clone(),
297                writer_tx.clone(),
298            )
299        });
300
301        let reconnect_max_attempts = None; // Stream mode does not reconnect
302        let connect_timeout = Duration::from_secs(10);
303
304        let reconnect_headers = ReconnectHeaders::new(std::mem::take(&mut config.headers));
305
306        Ok(Self {
307            config,
308            reconnect_headers,
309            handler: None, // Stream mode has no handler
310            ping_handler: None,
311            close_handler: None,
312            writer_tx,
313            connection_mode,
314            connection_epoch,
315            state_notify,
316            controller_notify,
317            reconnect_published,
318            connect_timeout,
319            heartbeat_timeout: None,
320            heartbeat_task,
321            read_task,
322            read_fence,
323            write_task,
324            backoff,
325            reconnect_throttle: ReconnectThrottle::default(),
326            reconnect_max_attempts,
327            reconnection_attempt_count: 0,
328            auth_tracker,
329            reconnect_buffer_waits_for_auth,
330            state_sink,
331            connection_rate_limit: None,
332        })
333    }
334
335    /// Creates an inner WebSocket client.
336    ///
337    /// # Errors
338    ///
339    /// Returns an error if:
340    /// - The writer capacity is invalid.
341    /// - The connection to the server fails.
342    /// - The exponential backoff configuration is invalid.
343    pub async fn connect_url(
344        config: WebSocketConfig,
345        message_handler: Option<MessageHandler>,
346        ping_handler: Option<PingHandler>,
347    ) -> Result<Self, TransportError> {
348        Self::connect_url_with_handler(
349            config,
350            message_handler.map(IncomingHandler::Message),
351            ping_handler.map(IncomingPingHandler::Ping),
352            None,
353            None,
354            None,
355            InitialConnectOptions::default(),
356        )
357        .await
358    }
359
360    async fn connect_url_with_handler(
361        config: WebSocketConfig,
362        handler: Option<IncomingHandler>,
363        ping_handler: Option<IncomingPingHandler>,
364        close_handler: Option<CloseHandler>,
365        state_sink: Option<SocketStateSink>,
366        connection_rate_limit: Option<ConnectionRateLimit>,
367        initial_connect_options: InitialConnectOptions,
368    ) -> Result<Self, TransportError> {
369        install_cryptographic_provider();
370
371        let is_stream_mode = handler.is_none();
372
373        // Adapters build this config by struct literal, bypassing the builder, so this is the only
374        // place the field invariants are enforced for them. Stream mode documents the reconnect and
375        // liveness fields as ignored and permits zero for them, so it checks only what it honors.
376        if is_stream_mode {
377            crate::writer::validate_capacity(config.writer_capacity).map_err(|e| {
378                TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
379            })?;
380
381            if config.heartbeat_interval_secs == Some(0) {
382                return Err(TransportError::Io(std::io::Error::new(
383                    std::io::ErrorKind::InvalidInput,
384                    "Heartbeat interval cannot be zero",
385                )));
386            }
387
388            validate_honored_inbound_limits(&config)?;
389        } else {
390            config.validate().map_err(|e| {
391                TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
392            })?;
393        }
394
395        let heartbeat_timeout = config.resolved_heartbeat_timeout().map(Duration::from_secs);
396        let reconnect_max_attempts = config.reconnect_max_attempts;
397
398        // Stream mode documents reconnect_* fields as ignored (callers may pass Some(0))
399        let connect_timeout = if is_stream_mode {
400            Duration::from_secs(10)
401        } else {
402            Duration::from_millis(config.connect_timeout_ms.unwrap_or(10_000))
403        };
404
405        let backoff = ExponentialBackoff::new(
406            Duration::from_millis(config.reconnect_delay_initial_ms.unwrap_or(2_000)),
407            Duration::from_millis(config.reconnect_delay_max_ms.unwrap_or(30_000)),
408            config.reconnect_backoff_factor.unwrap_or(1.5),
409            config.reconnect_jitter_ms.unwrap_or(100),
410            true, // immediate-first
411        )
412        .map_err(|e| {
413            TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
414        })?;
415
416        let cancellation_token = initial_connect_options
417            .cancellation_token
418            .unwrap_or_default();
419        let retry_policy = initial_connect_options.retry_policy;
420        let max_attempts = retry_policy
421            .as_ref()
422            .map_or(1, |policy| policy.max_attempts.get());
423        let retry_manager = RetryManager::new(initial_connect_retry_config(retry_policy.as_ref())?);
424        let attempt = AtomicU32::new(0);
425        let operation = || {
426            attempt.fetch_add(1, Ordering::Relaxed);
427
428            await_initial_connect_attempt(&cancellation_token, async {
429                if let Some(rate_limit) = &connection_rate_limit {
430                    rate_limit
431                        .limiter
432                        .await_keys_ready(Some(&rate_limit.keys))
433                        .await;
434                }
435
436                // Bound only the dial: the connection rate-limit wait has its own venue timing
437                dst::time::timeout(
438                    connect_timeout,
439                    Box::pin(Self::connect_with_server(
440                        &config.url,
441                        config.headers.clone(),
442                        config.backend,
443                        config.proxy_url.as_deref(),
444                        config.max_message_size_bytes,
445                        config.max_frame_size_bytes,
446                    )),
447                )
448                .await
449                .unwrap_or_else(|_| {
450                    Err(TransportError::Io(std::io::Error::new(
451                        std::io::ErrorKind::TimedOut,
452                        format!(
453                            "connection timed out after {}s",
454                            connect_timeout.as_secs_f64()
455                        ),
456                    )))
457                })
458            })
459        };
460
461        let classify = |error: &TransportError| {
462            let retryable = is_retryable_initial_connect_error(error);
463            let attempt = attempt.load(Ordering::Relaxed);
464            if retryable && attempt < max_attempts && !cancellation_token.is_cancelled() {
465                log::warn!(
466                    "WebSocket connection attempt {attempt}/{max_attempts} to {REDACTED} failed: {}",
467                    error.redacted_message()
468                );
469            }
470
471            retryable
472        };
473
474        let transport = retry_manager
475            .invocation(
476                INITIAL_CONNECT_OPERATION,
477                operation,
478                classify,
479                initial_connect_retry_error,
480            )
481            .cancellation_token(&cancellation_token)
482            .execute()
483            .await?;
484
485        let attempt = attempt.load(Ordering::Relaxed);
486        if attempt > 1 {
487            log::info!("WebSocket connection established after {attempt} attempts");
488        }
489
490        let (writer, reader) = transport;
491        let mut config = config;
492        let reconnect_headers = ReconnectHeaders::new(std::mem::take(&mut config.headers));
493
494        let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
495        let connection_epoch = Arc::new(AtomicU64::new(0));
496        let state_notify = Arc::new(tokio::sync::Notify::new());
497        let controller_notify = Arc::new(tokio::sync::Notify::new());
498        let reconnect_published = Arc::new(AtomicBool::new(true));
499        let outcome =
500            ConnectionMode::complete_reconnect_with_sink(&connection_mode, state_sink.as_ref());
501        debug_assert_eq!(outcome, ReconnectOutcome::Reconnected);
502
503        let (read_task, read_fence) = if is_stream_mode {
504            (None, None)
505        } else {
506            let read_fence = ReadSessionFence::new();
507            let read_task = Self::spawn_message_handler_task(
508                Arc::clone(&connection_mode),
509                Arc::clone(&state_notify),
510                read_fence.clone(),
511                reader,
512                0,
513                handler.as_ref(),
514                ping_handler.as_ref(),
515                close_handler.as_ref(),
516                state_sink.as_ref(),
517                config.idle_timeout_ms,
518                heartbeat_timeout,
519            );
520            (Some(read_task), Some(read_fence))
521        };
522
523        let auth_tracker = Arc::new(OnceLock::new());
524        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
525
526        let (writer_tx, writer_rx) = crate::writer::channel::<WriterCommand>(
527            config.writer_capacity.unwrap_or(DEFAULT_WRITER_CAPACITY),
528        );
529        let write_task = Self::spawn_write_task(
530            Arc::clone(&connection_mode),
531            Arc::clone(&controller_notify),
532            Arc::clone(&reconnect_published),
533            writer,
534            writer_rx,
535            Arc::clone(&connection_epoch),
536            Arc::clone(&auth_tracker),
537            Arc::clone(&reconnect_buffer_waits_for_auth),
538            state_sink.clone(),
539        );
540
541        // Optionally spawn a heartbeat task to periodically ping server
542        let heartbeat_task = config.heartbeat_interval_secs.map(|heartbeat_secs| {
543            Self::spawn_heartbeat_task(
544                Arc::clone(&connection_mode),
545                heartbeat_secs,
546                config.heartbeat_payload.clone(),
547                writer_tx.clone(),
548            )
549        });
550
551        Ok(Self {
552            config,
553            reconnect_headers,
554            handler,
555            ping_handler,
556            close_handler,
557            read_task,
558            read_fence,
559            write_task,
560            writer_tx,
561            heartbeat_task,
562            connection_mode,
563            connection_epoch,
564            state_notify,
565            controller_notify,
566            reconnect_published,
567            connect_timeout,
568            heartbeat_timeout,
569            backoff,
570            reconnect_throttle: ReconnectThrottle::default(),
571            reconnect_max_attempts,
572            reconnection_attempt_count: 0,
573            auth_tracker,
574            reconnect_buffer_waits_for_auth,
575            state_sink,
576            connection_rate_limit,
577        })
578    }
579
580    /// Connects to the server and returns the split halves of the active transport.
581    ///
582    /// Dispatches on `backend` to the matching transport implementation. The
583    /// [`TransportBackend::Tungstenite`] backend is always available; the
584    /// [`TransportBackend::Sockudo`] backend requires the `transport-sockudo`
585    /// Cargo feature (enabled by default) and uses a custom HTTP/1.1 handshake
586    /// path for upgrade headers.
587    ///
588    /// When `proxy_url` is `Some`, both backends establish an HTTP `CONNECT`
589    /// tunnel through the proxy before performing the WebSocket handshake, and
590    /// each keeps its own handshake path over the resulting stream.
591    ///
592    /// `max_message_size_bytes` and `max_frame_size_bytes` override the selected backend's inbound
593    /// caps. `None` passes that backend's current default config.
594    ///
595    /// # Errors
596    ///
597    /// Returns an error if the URL is invalid, headers fail to parse, the TCP / TLS layer cannot be
598    /// established, the proxy refuses the tunnel, or the WebSocket handshake is rejected by the
599    /// peer. Also returns an error if the Sockudo backend is selected without the
600    /// `transport-sockudo` feature.
601    #[inline]
602    pub async fn connect_with_server(
603        url: &str,
604        headers: Vec<(String, String)>,
605        backend: TransportBackend,
606        proxy_url: Option<&str>,
607        max_message_size_bytes: Option<usize>,
608        max_frame_size_bytes: Option<usize>,
609    ) -> Result<(MessageWriter, MessageReader), TransportError> {
610        #[cfg(all(feature = "simulation", madsim))]
611        if backend == TransportBackend::Sockudo {
612            return Err(TransportError::Other(
613                "Sockudo timers are unsupported under simulation; select Tungstenite".into(),
614            ));
615        }
616
617        #[cfg(all(feature = "simulation", madsim))]
618        if proxy_url.is_some() || !url.starts_with("ws://") {
619            return Err(TransportError::InvalidUrl(
620                "WebSocket simulation requires plaintext ws:// endpoints without a proxy".into(),
621            ));
622        }
623
624        let limits = InboundSizeLimits {
625            max_message_size_bytes,
626            max_frame_size_bytes,
627        };
628
629        match backend {
630            TransportBackend::Tungstenite => match proxy_url {
631                Some(proxy) => {
632                    Box::pin(Self::connect_tungstenite_via_proxy(
633                        url, headers, proxy, limits,
634                    ))
635                    .await
636                }
637                None => Self::connect_tungstenite(url, headers, limits).await,
638            },
639            TransportBackend::Sockudo => {
640                #[cfg(feature = "transport-sockudo")]
641                {
642                    match proxy_url {
643                        Some(proxy) => {
644                            Box::pin(Self::connect_sockudo_via_proxy(url, headers, proxy, limits))
645                                .await
646                        }
647                        None => Self::connect_sockudo(url, headers, limits).await,
648                    }
649                }
650
651                #[cfg(not(feature = "transport-sockudo"))]
652                {
653                    Err(TransportError::Other(
654                        "sockudo backend selected but the transport-sockudo \
655                         Cargo feature is not enabled"
656                            .to_string(),
657                    ))
658                }
659            }
660        }
661    }
662
663    /// Connects with the server creating a tokio-tungstenite websocket stream.
664    /// Production path using `connect_async_with_config`.
665    #[inline]
666    #[cfg(not(feature = "turmoil"))]
667    async fn connect_tungstenite(
668        url: &str,
669        headers: Vec<(String, String)>,
670        limits: InboundSizeLimits,
671    ) -> Result<(MessageWriter, MessageReader), TransportError> {
672        let request = tungstenite_request(url, headers)?;
673        let protocol = tungstenite_protocol_config(limits);
674
675        // Nagle stays enabled here so `apply_socket_options` owns every socket option in one place
676        #[cfg(not(all(feature = "simulation", madsim)))]
677        let (stream, _resp) = connect_async_with_config(request, protocol, false)
678            .await
679            .map_err(TransportError::from)?;
680        #[cfg(not(all(feature = "simulation", madsim)))]
681        crate::net::apply_socket_options(stream.get_ref().get_ref());
682
683        #[cfg(all(feature = "simulation", madsim))]
684        let (stream, _resp) = {
685            let target = request.uri();
686            let host = target
687                .host()
688                .ok_or_else(|| TransportError::InvalidUrl("missing hostname".into()))?;
689            let address = format!("{host}:{}", target.port_u16().unwrap_or(80));
690            let socket = dst::net::TcpStream::connect(address.as_str()).await?;
691            tokio_tungstenite::client_async_with_config(request, socket, protocol)
692                .await
693                .map_err(TransportError::from)?
694        };
695
696        let transport: BoxedWsTransport = Box::pin(TungsteniteTransport::new(stream));
697        Ok(transport.split())
698    }
699
700    /// Connects via an HTTP `CONNECT` proxy and performs the WebSocket
701    /// handshake over the resulting tunnel.
702    ///
703    /// Recognized but unsupported proxy schemes (currently SOCKS) log a
704    /// warning and fall back to a direct connection so existing REST proxy
705    /// configs remain usable. Only available in production builds; the
706    /// turmoil simulator does not model arbitrary outbound TCP via a proxy.
707    #[inline]
708    #[cfg(not(feature = "turmoil"))]
709    async fn connect_tungstenite_via_proxy(
710        url: &str,
711        headers: Vec<(String, String)>,
712        proxy_url: &str,
713        limits: InboundSizeLimits,
714    ) -> Result<(MessageWriter, MessageReader), TransportError> {
715        let proxy = match ProxyKind::parse(proxy_url)? {
716            ProxyKind::Http(target) => target,
717            ProxyKind::Unsupported { scheme } => {
718                log::warn!(
719                    "WebSocket proxy_url scheme '{scheme}' is not yet supported; \
720                     connecting without a WebSocket proxy"
721                );
722                return Self::connect_tungstenite(url, headers, limits).await;
723            }
724        };
725
726        let request = tungstenite_request(url, headers)?;
727
728        let target = WsTarget::parse(url)?;
729        let stream = tunnel_via_proxy(&target, &proxy).await?;
730
731        // `ProxiedStream` implements the IO traits over all four variants, so one
732        // instantiation covers every tunnel shape. The future is boxed because
733        // `client_async` produces a large state machine.
734        let transport: BoxedWsTransport =
735            Box::pin(proxied_ws_handshake(request, stream, limits)).await?;
736
737        Ok(transport.split())
738    }
739
740    /// Turmoil simulator variant: HTTP `CONNECT` tunneling is not supported
741    /// under the simulator so any proxy URL is rejected up front.
742    #[inline]
743    #[cfg(feature = "turmoil")]
744    #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
745    #[expect(
746        clippy::unused_async,
747        clippy::unused_async_trait_impl,
748        reason = "signature mirrors the production variant; both are awaited in the dispatcher"
749    )]
750    async fn connect_tungstenite_via_proxy(
751        _url: &str,
752        _headers: Vec<(String, String)>,
753        _proxy_url: &str,
754        _limits: InboundSizeLimits,
755    ) -> Result<(MessageWriter, MessageReader), TransportError> {
756        Err(TransportError::Other(
757            "proxy_url is not supported under the turmoil simulator".to_string(),
758        ))
759    }
760
761    /// Connects with the server creating a tokio-tungstenite websocket stream.
762    /// Turmoil version that uses the lower-level `client_async` API with injected stream.
763    #[inline]
764    #[cfg(feature = "turmoil")]
765    async fn connect_tungstenite(
766        url: &str,
767        headers: Vec<(String, String)>,
768        limits: InboundSizeLimits,
769    ) -> Result<(MessageWriter, MessageReader), TransportError> {
770        let request = tungstenite_request(url, headers)?;
771
772        let uri = request.uri();
773        let scheme = uri.scheme_str().unwrap_or("ws");
774        let host = uri
775            .host()
776            .ok_or_else(|| TransportError::InvalidUrl("missing hostname".to_string()))?;
777
778        // Determine port: use explicit port if specified, otherwise default based on scheme
779        let port = uri
780            .port_u16()
781            .unwrap_or_else(|| if scheme == "wss" { 443 } else { 80 });
782
783        let addr = format!("{host}:{port}");
784
785        // Use the connector to get a turmoil-compatible stream
786        let connector = crate::net::RealTcpConnector;
787        let tcp_stream = connector.connect(&addr).await?;
788        crate::net::apply_socket_options(&tcp_stream);
789
790        // Wrap stream appropriately based on scheme
791        let maybe_tls_stream = if scheme == "wss" {
792            // Build TLS config with webpki roots
793            let mut root_store = rustls::RootCertStore::empty();
794            root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
795
796            let config = ClientConfig::builder()
797                .with_root_certificates(root_store)
798                .with_no_client_auth();
799
800            let tls_connector = TlsConnector::from(std::sync::Arc::new(config));
801            let domain = rustls::pki_types::ServerName::try_from(host.to_string())
802                .map_err(|e| TransportError::Tls(format!("Invalid DNS name: {e}")))?;
803
804            let tls_stream = tls_connector
805                .connect(domain, tcp_stream)
806                .await
807                .map_err(TransportError::Io)?;
808            MaybeTlsStream::Rustls(tls_stream)
809        } else {
810            MaybeTlsStream::Plain(tcp_stream)
811        };
812
813        // Use client_async with the stream (plain or TLS)
814        let (stream, _resp) = client_async_with_config(
815            request,
816            maybe_tls_stream,
817            tungstenite_protocol_config(limits),
818        )
819        .await
820        .map_err(TransportError::from)?;
821        let transport: BoxedWsTransport = Box::pin(TungsteniteTransport::new(stream));
822        Ok(transport.split())
823    }
824
825    /// Connects with the server using the sockudo-ws backend.
826    ///
827    /// Uses a local HTTP/1.1 handshake path so error logging and stream
828    /// construction stay in our hands regardless of header count.
829    ///
830    /// Under the turmoil simulator, only plaintext `ws://` is supported (the
831    /// simulator does not model TLS), so a `wss://` URL returns
832    /// [`TransportError::Tls`] up front.
833    #[inline]
834    #[cfg(feature = "transport-sockudo")]
835    async fn connect_sockudo(
836        url: &str,
837        headers: Vec<(String, String)>,
838        limits: InboundSizeLimits,
839    ) -> Result<(MessageWriter, MessageReader), TransportError> {
840        let target = SockudoTarget::parse(url)?;
841        validate_extra_headers(&headers).map_err(TransportError::from)?;
842
843        #[cfg(feature = "turmoil")]
844        if target.is_tls {
845            return Err(TransportError::Tls(
846                "wss:// is not supported under the turmoil simulator; use ws://".to_string(),
847            ));
848        }
849
850        let tcp_stream = TcpStream::connect((target.host.as_str(), target.port))
851            .await
852            .map_err(TransportError::Io)?;
853
854        crate::net::apply_socket_options(&tcp_stream);
855
856        #[cfg(not(feature = "turmoil"))]
857        if target.is_tls {
858            let mut root_store = rustls::RootCertStore::empty();
859            root_store.extend(webpki_roots::TLS_SERVER_ROOTS.iter().cloned());
860            let config = ClientConfig::builder()
861                .with_root_certificates(root_store)
862                .with_no_client_auth();
863            let connector = TlsConnector::from(std::sync::Arc::new(config));
864            let domain = rustls::pki_types::ServerName::try_from(target.host.clone())
865                .map_err(|e| TransportError::Tls(format!("Invalid DNS name: {e}")))?;
866            let tls_stream = connector
867                .connect(domain, tcp_stream)
868                .await
869                .map_err(TransportError::Io)?;
870            return Self::finish_sockudo_handshake(tls_stream, &target, &headers, limits).await;
871        }
872
873        Self::finish_sockudo_handshake(tcp_stream, &target, &headers, limits).await
874    }
875
876    /// Connects via an HTTP `CONNECT` proxy and performs the sockudo WebSocket
877    /// handshake over the resulting tunnel.
878    ///
879    /// Recognized but unsupported proxy schemes (currently SOCKS) log a warning
880    /// and fall back to a direct connection, matching the Tungstenite path.
881    #[inline]
882    #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
883    async fn connect_sockudo_via_proxy(
884        url: &str,
885        headers: Vec<(String, String)>,
886        proxy_url: &str,
887        limits: InboundSizeLimits,
888    ) -> Result<(MessageWriter, MessageReader), TransportError> {
889        let proxy = match ProxyKind::parse(proxy_url)? {
890            ProxyKind::Http(target) => target,
891            ProxyKind::Unsupported { scheme } => {
892                log::warn!(
893                    "WebSocket proxy_url scheme '{scheme}' is not yet supported; \
894                     connecting without a WebSocket proxy"
895                );
896                return Self::connect_sockudo(url, headers, limits).await;
897            }
898        };
899
900        let target = SockudoTarget::parse(url)?;
901        validate_extra_headers(&headers).map_err(TransportError::from)?;
902
903        // `tunnel_via_proxy` establishes upstream TLS inside the tunnel when the
904        // target is `wss://`, so the handshake below runs over the finished stream
905        // regardless of which of the four tunnel shapes it returned.
906        let ws_target = WsTarget::parse(url)?;
907        let stream = tunnel_via_proxy(&ws_target, &proxy).await?;
908
909        Self::finish_sockudo_handshake(stream, &target, &headers, limits).await
910    }
911
912    /// Turmoil simulator variant: HTTP `CONNECT` tunneling is not modeled under
913    /// the simulator so any proxy URL is rejected up front.
914    #[inline]
915    #[cfg(all(feature = "transport-sockudo", feature = "turmoil"))]
916    #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
917    #[expect(
918        clippy::unused_async,
919        clippy::unused_async_trait_impl,
920        reason = "signature mirrors the production variant; both are awaited in the dispatcher"
921    )]
922    async fn connect_sockudo_via_proxy(
923        _url: &str,
924        _headers: Vec<(String, String)>,
925        _proxy_url: &str,
926        _limits: InboundSizeLimits,
927    ) -> Result<(MessageWriter, MessageReader), TransportError> {
928        Err(TransportError::Other(
929            "proxy_url is not supported under the turmoil simulator".to_string(),
930        ))
931    }
932
933    #[cfg(feature = "transport-sockudo")]
934    async fn finish_sockudo_handshake<S>(
935        mut stream: S,
936        target: &SockudoTarget,
937        headers: &[(String, String)],
938        limits: InboundSizeLimits,
939    ) -> Result<(MessageWriter, MessageReader), TransportError>
940    where
941        S: AsyncRead + AsyncWrite + Unpin + Send + 'static,
942    {
943        // Use one path for uniform error logging and ownership of stream
944        // construction; sockudo's handshake error drops the rejection status
945        // the retry policy needs.
946        let handshake =
947            client_handshake_with_headers(&mut stream, &target.host_header, &target.path, headers)
948                .await?;
949
950        // Reading the HTTP 101 may also read the first WebSocket frame prefix;
951        // replay it only when present so the ordinary path stays unwrapped.
952        let stream = match handshake.leftover {
953            Some(prefix) => SockudoStream::<Http1>::new(PrefixedIo::new(stream, prefix)),
954            None => SockudoStream::<Http1>::new(stream),
955        };
956
957        let ws =
958            SockudoWebSocketStream::from_raw(stream, Role::Client, sockudo_protocol_config(limits));
959        let transport: BoxedWsTransport = Box::pin(SockudoTransport::new(ws));
960        Ok(transport.split())
961    }
962}
963
964#[derive(Clone, Copy, Debug, Default)]
965struct InboundSizeLimits {
966    max_message_size_bytes: Option<usize>,
967    max_frame_size_bytes: Option<usize>,
968}
969
970fn validate_honored_inbound_limits(config: &WebSocketConfig) -> Result<(), TransportError> {
971    if config.max_message_size_bytes == Some(0) {
972        return Err(TransportError::Io(std::io::Error::new(
973            std::io::ErrorKind::InvalidInput,
974            "max_message_size_bytes must be positive",
975        )));
976    }
977
978    if config.max_frame_size_bytes == Some(0) {
979        return Err(TransportError::Io(std::io::Error::new(
980            std::io::ErrorKind::InvalidInput,
981            "max_frame_size_bytes must be positive",
982        )));
983    }
984
985    Ok(())
986}
987
988fn tungstenite_protocol_config(
989    limits: InboundSizeLimits,
990) -> Option<tokio_tungstenite::tungstenite::protocol::WebSocketConfig> {
991    if limits.max_message_size_bytes.is_none() && limits.max_frame_size_bytes.is_none() {
992        return None;
993    }
994
995    let mut config = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default();
996    if let Some(size) = limits.max_message_size_bytes {
997        config.max_message_size = Some(size);
998    }
999
1000    if let Some(size) = limits.max_frame_size_bytes {
1001        config.max_frame_size = Some(size);
1002    }
1003
1004    Some(config)
1005}
1006
1007#[cfg(feature = "transport-sockudo")]
1008fn sockudo_protocol_config(limits: InboundSizeLimits) -> SockudoConfig {
1009    let mut config = SockudoConfig::default();
1010    if let Some(size) = limits.max_message_size_bytes {
1011        config.max_message_size = size;
1012    }
1013
1014    if let Some(size) = limits.max_frame_size_bytes {
1015        config.max_frame_size = size;
1016    }
1017
1018    config
1019}
1020
1021fn tungstenite_request(
1022    url: &str,
1023    headers: Vec<(String, String)>,
1024) -> Result<Request, TransportError> {
1025    let mut request = url.into_client_request().map_err(TransportError::from)?;
1026
1027    for (key, value) in headers {
1028        let value = HeaderValue::from_str(&value)
1029            .map_err(|e| TransportError::Handshake(format!("invalid header value: {e}")))?;
1030        let name: HeaderName = key
1031            .parse()
1032            .map_err(|e| TransportError::Handshake(format!("invalid header name: {e}")))?;
1033        request.headers_mut().insert(name, value);
1034    }
1035
1036    Ok(request)
1037}
1038
1039fn is_connection_drop_transport_error(err: &TransportError) -> bool {
1040    err.is_closed() || matches!(err, TransportError::Io(e) if is_connection_drop_io_error(e))
1041}
1042
1043fn is_retryable_initial_connect_error(err: &TransportError) -> bool {
1044    match err {
1045        TransportError::ConnectionClosed
1046        | TransportError::ConnectionReset
1047        | TransportError::ClosedByPeer(_) => true,
1048        TransportError::Io(error) => !matches!(
1049            error.kind(),
1050            std::io::ErrorKind::InvalidInput
1051                | std::io::ErrorKind::InvalidData
1052                | std::io::ErrorKind::Unsupported
1053                | std::io::ErrorKind::PermissionDenied
1054        ),
1055        TransportError::UpgradeRejected(status) | TransportError::ProxyConnectRejected(status) => {
1056            retryable_status(*status)
1057        }
1058        TransportError::InvalidUrl(_)
1059        | TransportError::Handshake(_)
1060        | TransportError::Tls(_)
1061        | TransportError::Protocol(_)
1062        | TransportError::MessageTooLarge
1063        | TransportError::FrameTooLarge
1064        | TransportError::InvalidUtf8
1065        | TransportError::Other(_) => false,
1066    }
1067}
1068
1069fn initial_connect_retry_config(
1070    policy: Option<&InitialConnectRetryPolicy>,
1071) -> Result<RetryConfig, TransportError> {
1072    if let Some(policy) = policy {
1073        ExponentialBackoff::new(
1074            policy.delay_initial,
1075            policy.delay_max,
1076            policy.backoff_factor,
1077            policy.jitter_ms,
1078            false,
1079        )
1080        .map_err(|e| {
1081            TransportError::Io(std::io::Error::new(std::io::ErrorKind::InvalidInput, e))
1082        })?;
1083
1084        Ok(RetryConfig {
1085            max_retries: policy.max_attempts.get() - 1,
1086            initial_delay_ms: duration_to_millis("delay_initial", policy.delay_initial)?,
1087            max_delay_ms: duration_to_millis("delay_max", policy.delay_max)?,
1088            backoff_factor: policy.backoff_factor,
1089            jitter_ms: policy.jitter_ms,
1090            operation_timeout_ms: None,
1091            immediate_first: false,
1092            max_elapsed_ms: None,
1093        })
1094    } else {
1095        Ok(RetryConfig {
1096            max_retries: 0,
1097            initial_delay_ms: 1,
1098            max_delay_ms: 1,
1099            backoff_factor: 1.0,
1100            jitter_ms: 0,
1101            operation_timeout_ms: None,
1102            immediate_first: false,
1103            max_elapsed_ms: None,
1104        })
1105    }
1106}
1107
1108fn duration_to_millis(field: &str, duration: Duration) -> Result<u64, TransportError> {
1109    let nanoseconds = duration.as_nanos();
1110    if nanoseconds > u128::from(u64::MAX) {
1111        return Err(TransportError::Io(std::io::Error::new(
1112            std::io::ErrorKind::InvalidInput,
1113            format!("{field} exceeds the maximum backoff duration"),
1114        )));
1115    }
1116
1117    let milliseconds = nanoseconds
1118        .div_ceil(1_000_000)
1119        .min(u128::from(u64::MAX / 1_000_000));
1120    u64::try_from(milliseconds).map_err(|_| {
1121        TransportError::Io(std::io::Error::new(
1122            std::io::ErrorKind::InvalidInput,
1123            format!("{field} exceeds the maximum backoff duration"),
1124        ))
1125    })
1126}
1127
1128async fn await_initial_connect_attempt<F, T>(
1129    cancellation_token: &CancellationToken,
1130    attempt: F,
1131) -> Result<T, TransportError>
1132where
1133    F: Future<Output = Result<T, TransportError>>,
1134{
1135    tokio::select! {
1136        biased;
1137        () = cancellation_token.cancelled() => Err(initial_connect_cancelled()),
1138        result = attempt => result,
1139    }
1140}
1141
1142fn initial_connect_retry_error(error: RetryError) -> TransportError {
1143    let kind = match error {
1144        RetryError::Canceled => return initial_connect_cancelled(),
1145        RetryError::InvalidConfiguration { .. } => std::io::ErrorKind::InvalidInput,
1146        RetryError::OperationTimeout { .. } | RetryError::ElapsedBudgetExceeded { .. } => {
1147            std::io::ErrorKind::TimedOut
1148        }
1149    };
1150
1151    TransportError::Io(std::io::Error::new(kind, error))
1152}
1153
1154fn initial_connect_cancelled() -> TransportError {
1155    TransportError::Io(std::io::Error::new(
1156        std::io::ErrorKind::Interrupted,
1157        "initial WebSocket connection cancelled",
1158    ))
1159}
1160
1161// Debug when we asked to disconnect (Disconnect/Closed), else Warn for a peer close
1162fn read_termination_log_level(connection_state: &AtomicU8) -> log::Level {
1163    let mode = ConnectionMode::from_atomic(connection_state);
1164    if mode.is_disconnect() || mode.is_closed() {
1165        log::Level::Debug
1166    } else {
1167        log::Level::Warn
1168    }
1169}
1170
1171#[cfg(test)]
1172mod connection_error_tests {
1173    use std::io;
1174
1175    use rstest::rstest;
1176
1177    use super::*;
1178    use crate::transport::CloseFrame;
1179
1180    #[rstest]
1181    #[case(TransportError::ConnectionClosed, true)]
1182    #[case(TransportError::ConnectionReset, true)]
1183    #[case(TransportError::ClosedByPeer(Some(CloseFrame::new(1000, "bye"))), true)]
1184    #[case(TransportError::ClosedByPeer(None), true)]
1185    #[case(TransportError::Io(io::Error::from(io::ErrorKind::BrokenPipe)), true)]
1186    #[case(
1187        TransportError::Io(io::Error::from(io::ErrorKind::ConnectionReset)),
1188        true
1189    )]
1190    #[case(TransportError::Io(io::Error::from(io::ErrorKind::TimedOut)), true)]
1191    #[case(
1192        TransportError::Io(io::Error::from(io::ErrorKind::UnexpectedEof)),
1193        true
1194    )]
1195    #[case(
1196        TransportError::Io(io::Error::from(io::ErrorKind::InvalidInput)),
1197        false
1198    )]
1199    #[case(TransportError::InvalidUrl("http://example.com".into()), false)]
1200    #[case(TransportError::Handshake("bad".into()), false)]
1201    #[case(TransportError::Protocol("bad opcode".into()), false)]
1202    #[case(TransportError::Tls("bad certificate".into()), false)]
1203    #[case(TransportError::MessageTooLarge, false)]
1204    #[case(TransportError::FrameTooLarge, false)]
1205    #[case(TransportError::InvalidUtf8, false)]
1206    #[case(TransportError::Other("backend protocol mismatch".into()), false)]
1207    fn connection_drop_transport_error_classification(
1208        #[case] err: TransportError,
1209        #[case] expected: bool,
1210    ) {
1211        assert_eq!(is_connection_drop_transport_error(&err), expected);
1212    }
1213
1214    #[rstest]
1215    #[case(Duration::ZERO, 0)]
1216    #[case(Duration::from_nanos(1), 1)]
1217    #[case(Duration::from_micros(500), 1)]
1218    #[case(Duration::from_nanos(1_000_001), 2)]
1219    #[case(Duration::from_millis(500), 500)]
1220    #[case(Duration::from_secs(5), 5_000)]
1221    fn duration_to_millis_rounds_up(#[case] duration: Duration, #[case] expected: u64) {
1222        assert_eq!(duration_to_millis("delay", duration).unwrap(), expected);
1223    }
1224
1225    #[rstest]
1226    fn duration_to_millis_caps_at_backoff_maximum() {
1227        assert_eq!(
1228            duration_to_millis("delay", Duration::from_nanos(u64::MAX)).unwrap(),
1229            u64::MAX / 1_000_000
1230        );
1231    }
1232
1233    #[rstest]
1234    fn duration_to_millis_rejects_excessive_duration() {
1235        let error = duration_to_millis("delay", Duration::from_secs(u64::MAX)).unwrap_err();
1236
1237        let TransportError::Io(error) = error else {
1238            panic!("expected an I/O error, was: {error:?}");
1239        };
1240
1241        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
1242        assert_eq!(
1243            error.to_string(),
1244            "delay exceeds the maximum backoff duration"
1245        );
1246    }
1247
1248    #[rstest]
1249    fn initial_connect_retry_config_preserves_duration_validation() {
1250        let policy = InitialConnectRetryPolicy {
1251            max_attempts: std::num::NonZeroU32::new(2).unwrap(),
1252            delay_initial: Duration::from_micros(900),
1253            delay_max: Duration::from_micros(500),
1254            backoff_factor: 2.0,
1255            jitter_ms: 0,
1256        };
1257
1258        let error = initial_connect_retry_config(Some(&policy)).unwrap_err();
1259
1260        let TransportError::Io(error) = error else {
1261            panic!("expected an I/O error, was: {error:?}");
1262        };
1263
1264        assert_eq!(error.kind(), io::ErrorKind::InvalidInput);
1265        assert!(
1266            error
1267                .to_string()
1268                .contains("delay_max must be >= delay_initial")
1269        );
1270    }
1271
1272    #[rstest]
1273    #[tokio::test]
1274    async fn initial_connect_attempt_prefers_cancellation_over_ready_result() {
1275        let token = CancellationToken::new();
1276        token.cancel();
1277
1278        let error = await_initial_connect_attempt(&token, std::future::ready(Ok(())))
1279            .await
1280            .expect_err("cancellation should take priority");
1281
1282        assert!(
1283            matches!(error, TransportError::Io(ref error) if error.kind() == io::ErrorKind::Interrupted)
1284        );
1285    }
1286
1287    #[rstest]
1288    fn unset_inbound_limits_preserve_backend_defaults() {
1289        let limits = InboundSizeLimits::default();
1290        assert!(tungstenite_protocol_config(limits).is_none());
1291
1292        #[cfg(feature = "transport-sockudo")]
1293        {
1294            let config = sockudo_protocol_config(limits);
1295            let default = SockudoConfig::default();
1296            assert_eq!(config.max_message_size, default.max_message_size);
1297            assert_eq!(config.max_frame_size, default.max_frame_size);
1298        }
1299    }
1300
1301    #[rstest]
1302    fn configured_inbound_limit_overrides_only_that_cap() {
1303        let limits = InboundSizeLimits {
1304            max_message_size_bytes: Some(32),
1305            max_frame_size_bytes: None,
1306        };
1307
1308        let config = tungstenite_protocol_config(limits).expect("override should build a config");
1309        let default = tokio_tungstenite::tungstenite::protocol::WebSocketConfig::default();
1310
1311        assert_eq!(config.max_message_size, Some(32));
1312        assert_eq!(config.max_frame_size, default.max_frame_size);
1313
1314        #[cfg(feature = "transport-sockudo")]
1315        {
1316            let config = sockudo_protocol_config(limits);
1317            let default = SockudoConfig::default();
1318            assert_eq!(config.max_message_size, 32);
1319            assert_eq!(config.max_frame_size, default.max_frame_size);
1320        }
1321    }
1322}
1323
1324/// Complete the WebSocket handshake over a stream that has already been
1325/// tunneled through an HTTP `CONNECT` proxy. Generic over the concrete
1326/// stream type so the four [`super::proxy::ProxiedStream`] variants share
1327/// a single body.
1328#[cfg(not(feature = "turmoil"))]
1329async fn proxied_ws_handshake<S>(
1330    request: tokio_tungstenite::tungstenite::handshake::client::Request,
1331    stream: S,
1332    limits: InboundSizeLimits,
1333) -> Result<BoxedWsTransport, TransportError>
1334where
1335    S: tokio::io::AsyncRead + tokio::io::AsyncWrite + Unpin + Send + 'static,
1336{
1337    let (ws, _resp) = tokio_tungstenite::client_async_with_config(
1338        request,
1339        stream,
1340        tungstenite_protocol_config(limits),
1341    )
1342    .await
1343    .map_err(TransportError::from)?;
1344    Ok(Box::pin(TungsteniteTransport::new(ws)))
1345}
1346
1347/// Parsed components of a `ws://` / `wss://` URL needed by the sockudo backend.
1348///
1349/// Sockudo's HTTP/1.1 client passes the `host` argument verbatim as the
1350/// HTTP `Host:` header, so it must include the explicit port when one is
1351/// present in the URL (RFC 7230 section 5.4). The DNS / SNI lookup uses the bare
1352/// host without the port.
1353#[cfg(feature = "transport-sockudo")]
1354#[derive(Debug, PartialEq, Eq)]
1355struct SockudoTarget {
1356    host: String,
1357    host_header: String,
1358    port: u16,
1359    path: String,
1360    is_tls: bool,
1361}
1362
1363#[cfg(feature = "transport-sockudo")]
1364impl SockudoTarget {
1365    fn parse(url: &str) -> Result<Self, TransportError> {
1366        let parsed = url::Url::parse(url)
1367            .map_err(|e| TransportError::InvalidUrl(format!("invalid WebSocket URL: {e}")))?;
1368
1369        let scheme = parsed.scheme();
1370
1371        let is_tls = match scheme {
1372            "ws" => false,
1373            "wss" => true,
1374            other => {
1375                return Err(TransportError::InvalidUrl(format!(
1376                    "expected ws:// or wss:// scheme, was {other}"
1377                )));
1378            }
1379        };
1380
1381        let raw_host = parsed
1382            .host_str()
1383            .ok_or_else(|| TransportError::InvalidUrl("missing hostname".to_string()))?;
1384
1385        // url::Url stores IPv6 hosts in their bracketed form (e.g. `[::1]`).
1386        // Brackets are correct for the HTTP `Host:` header but invalid for
1387        // DNS/TCP and TLS SNI, so we keep two representations: a bracketed
1388        // `host_header` for the upgrade, and a bare `host` for the socket connection.
1389        let is_bracketed = raw_host.starts_with('[') && raw_host.ends_with(']');
1390        let host = if is_bracketed {
1391            raw_host[1..raw_host.len() - 1].to_string()
1392        } else {
1393            raw_host.to_string()
1394        };
1395
1396        let explicit_port = parsed.port();
1397        let port = explicit_port.unwrap_or(if is_tls { 443 } else { 80 });
1398
1399        let host_header = match explicit_port {
1400            Some(p) => format!("{raw_host}:{p}"),
1401            None => raw_host.to_string(),
1402        };
1403
1404        let path = if parsed.path().is_empty() {
1405            "/".to_string()
1406        } else {
1407            let mut p = parsed.path().to_string();
1408            if let Some(query) = parsed.query() {
1409                p.push('?');
1410                p.push_str(query);
1411            }
1412
1413            p
1414        };
1415
1416        Ok(Self {
1417            host,
1418            host_header,
1419            port,
1420            path,
1421            is_tls,
1422        })
1423    }
1424}
1425
1426impl WebSocketClientInner {
1427    /// Reconnect with server.
1428    ///
1429    /// Make a new connection with server. Use the new read and write halves
1430    /// to update self writer and read and heartbeat tasks.
1431    ///
1432    /// For stream-based clients (created via [`WebSocketClient::stream_builder`]), reconnection is
1433    /// disabled because the reader is owned by the caller and cannot be replaced. Stream users
1434    /// should handle disconnections by creating a new connection.
1435    ///
1436    /// The reconnect timeout bounds only connection establishment. Once the
1437    /// new writer is handed to the writer task the swap runs to completion,
1438    /// so buffered messages can never drain into a connection that lost its
1439    /// reader to a timeout; the post-connect steps are individually bounded
1440    /// by the writer task's graceful-shutdown timeout.
1441    ///
1442    /// # Errors
1443    ///
1444    /// Returns an error if:
1445    /// - The reconnection attempt times out.
1446    /// - The connection to the server fails.
1447    pub async fn reconnect(&mut self) -> Result<(), TransportError> {
1448        Box::pin(self.reconnect_with_outcome()).await.map(|_| ())
1449    }
1450
1451    async fn wait_for_reconnect_publication(&self) -> bool {
1452        let fallback_interval = Duration::from_millis(CONTROLLER_FALLBACK_INTERVAL_MS);
1453
1454        loop {
1455            let mut notified = pin!(self.controller_notify.notified());
1456            notified.as_mut().enable();
1457
1458            if !ConnectionMode::from_atomic(&self.connection_mode).is_reconnect() {
1459                return false;
1460            }
1461
1462            if self.reconnect_published.load(Ordering::SeqCst) {
1463                return true;
1464            }
1465
1466            tokio::select! {
1467                biased;
1468                () = notified => {}
1469                () = dst::time::sleep(fallback_interval) => {}
1470            }
1471        }
1472    }
1473
1474    async fn reconnect_with_outcome(&mut self) -> Result<ReconnectOutcome, TransportError> {
1475        log::info!("Reconnecting");
1476
1477        if self.handler.is_none() {
1478            log::warn!(
1479                "Auto-reconnect disabled for stream-based WebSocket client; \
1480                stream users must manually reconnect by creating a new connection"
1481            );
1482            // Transition to CLOSED state to stop reconnection attempts
1483            self.connection_mode
1484                .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
1485            fail_registered_auth(
1486                self.auth_tracker.as_ref(),
1487                "WebSocket stream mode cannot reconnect",
1488            );
1489            // Publish the terminal state only once its bookkeeping is complete; the
1490            // controller used to notify after `reconnect` returned, so a waiter must
1491            // not observe `Closed` before pending authentication has been failed.
1492            self.state_notify.notify_waiters();
1493            return Ok(ReconnectOutcome::Aborted);
1494        }
1495
1496        if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
1497            log::debug!("Reconnect aborted due to disconnect state");
1498            return Ok(ReconnectOutcome::Aborted);
1499        }
1500
1501        if let Some(rate_limit) = &self.connection_rate_limit {
1502            rate_limit
1503                .limiter
1504                .await_keys_ready(Some(&rate_limit.keys))
1505                .await;
1506        }
1507
1508        if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
1509            log::debug!("Reconnect aborted during connection rate-limit wait");
1510            return Ok(ReconnectOutcome::Aborted);
1511        }
1512
1513        // Bound only connection establishment; the swap below must run to completion
1514        let (new_writer, reader) = dst::time::timeout(
1515            self.connect_timeout,
1516            Box::pin(Self::connect_with_server(
1517                &self.config.url,
1518                self.reconnect_headers.snapshot(),
1519                self.config.backend,
1520                self.config.proxy_url.as_deref(),
1521                self.config.max_message_size_bytes,
1522                self.config.max_frame_size_bytes,
1523            )),
1524        )
1525        .await
1526        .map_err(|_| {
1527            TransportError::Io(std::io::Error::new(
1528                std::io::ErrorKind::TimedOut,
1529                format!(
1530                    "reconnection timed out after {}s",
1531                    self.connect_timeout.as_secs_f64()
1532                ),
1533            ))
1534        })??;
1535
1536        if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
1537            log::debug!("Reconnect aborted mid-flight (after connect)");
1538            return Ok(ReconnectOutcome::Aborted);
1539        }
1540
1541        // Use a oneshot channel to synchronize the writer swap before transitioning
1542        // back to ACTIVE. Buffered messages stay in the writer task and replay later.
1543        let (tx, rx) = tokio::sync::oneshot::channel();
1544
1545        if let Err(e) = self
1546            .writer_tx
1547            .send_update(WriterCommand::Update(new_writer, tx))
1548        {
1549            log::error!("{e}");
1550            return Err(TransportError::Io(std::io::Error::new(
1551                std::io::ErrorKind::BrokenPipe,
1552                format!("Failed to send update command: {e}"),
1553            )));
1554        }
1555
1556        // Wait for writer to confirm it accepted the new socket
1557        let connection_epoch = match rx.await {
1558            Ok(connection_epoch) => {
1559                log::debug!("Writer confirmed socket update: epoch={connection_epoch}");
1560                connection_epoch
1561            }
1562            Err(e) => {
1563                log::error!("Writer dropped update channel: {e}");
1564                return Err(TransportError::Io(std::io::Error::new(
1565                    std::io::ErrorKind::BrokenPipe,
1566                    "Writer task dropped response channel",
1567                )));
1568            }
1569        };
1570
1571        // Delay before closing connection
1572        dst::time::sleep(Duration::from_millis(GRACEFUL_SHUTDOWN_DELAY_MS)).await;
1573
1574        if ConnectionMode::from_atomic(&self.connection_mode).is_disconnect() {
1575            log::debug!("Reconnect aborted mid-flight (after delay)");
1576            return Ok(ReconnectOutcome::Aborted);
1577        }
1578
1579        if let Some(read_fence) = self.read_fence.take() {
1580            read_fence.invalidate();
1581        }
1582
1583        if let Some(ref read_task) = self.read_task.take()
1584            && !read_task.is_finished()
1585        {
1586            read_task.abort();
1587            log_task_aborted("read");
1588        }
1589
1590        if !self.wait_for_reconnect_publication().await {
1591            log::debug!("Reconnect aborted before state publication completed");
1592            return Ok(ReconnectOutcome::Aborted);
1593        }
1594
1595        // Atomically transition from Reconnect to Active
1596        // This prevents race condition where disconnect could be requested between check and store
1597        if ConnectionMode::complete_reconnect_with_sink(
1598            &self.connection_mode,
1599            self.state_sink.as_ref(),
1600        ) == ReconnectOutcome::Aborted
1601        {
1602            log::debug!("Reconnect aborted (state changed during reconnect)");
1603            return Ok(ReconnectOutcome::Aborted);
1604        }
1605
1606        let read_fence = ReadSessionFence::new();
1607        self.read_task = Some(Self::spawn_message_handler_task(
1608            Arc::clone(&self.connection_mode),
1609            Arc::clone(&self.state_notify),
1610            read_fence.clone(),
1611            reader,
1612            connection_epoch,
1613            self.handler.as_ref(),
1614            self.ping_handler.as_ref(),
1615            self.close_handler.as_ref(),
1616            self.state_sink.as_ref(),
1617            self.config.idle_timeout_ms,
1618            self.heartbeat_timeout,
1619        ));
1620        self.read_fence = Some(read_fence);
1621
1622        log::info!("Reconnect succeeded");
1623        Ok(ReconnectOutcome::Reconnected)
1624    }
1625
1626    /// Returns whether the client's transport tasks are still running.
1627    ///
1628    /// Returns `true` if both the read and write tasks are still running.
1629    /// There may be some delay between the connection closing and the
1630    /// client detecting it.
1631    #[inline]
1632    #[must_use]
1633    pub fn is_alive(&self) -> bool {
1634        match &self.read_task {
1635            Some(read_task) => !read_task.is_finished() && !self.write_task.is_finished(),
1636            None => !self.write_task.is_finished(),
1637        }
1638    }
1639
1640    #[expect(
1641        clippy::too_many_arguments,
1642        reason = "both handler modes share the same reader lifecycle"
1643    )]
1644    fn spawn_message_handler_task(
1645        connection_state: Arc<AtomicU8>,
1646        state_notify: Arc<tokio::sync::Notify>,
1647        read_fence: ReadSessionFence,
1648        mut reader: MessageReader,
1649        connection_epoch: u64,
1650        handler: Option<&IncomingHandler>,
1651        ping_handler: Option<&IncomingPingHandler>,
1652        close_handler: Option<&CloseHandler>,
1653        state_sink: Option<&SocketStateSink>,
1654        idle_timeout_ms: Option<u64>,
1655        heartbeat_timeout: Option<Duration>,
1656    ) -> dst::task::JoinHandle<()> {
1657        log::debug!("Started message handler task 'read'");
1658
1659        let check_interval = Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS);
1660        let idle_timeout = idle_timeout_ms.map(Duration::from_millis);
1661
1662        let handler = handler.cloned();
1663        let ping_handler = ping_handler.cloned();
1664        let close_handler = close_handler.cloned();
1665        let state_sink = state_sink.cloned();
1666
1667        dst::task::spawn(async move {
1668            let mut last_data_time = dst::time::Instant::now();
1669            let mut last_frame_time = dst::time::Instant::now();
1670
1671            loop {
1672                if !ConnectionMode::from_atomic(&connection_state).is_active()
1673                    || !read_fence.is_valid()
1674                {
1675                    break;
1676                }
1677
1678                let read_result = dst::time::timeout(check_interval, reader.next()).await;
1679
1680                if let Ok(Some(Ok(ref message))) = read_result
1681                    && (!ConnectionMode::from_atomic(&connection_state).is_active()
1682                        || !read_fence.is_valid())
1683                {
1684                    log::debug!(
1685                        "Dropping WebSocket message with {} bytes after session ended",
1686                        message.as_bytes().len()
1687                    );
1688                    break;
1689                }
1690
1691                if matches!(&read_result, Ok(Some(Ok(_)))) {
1692                    last_frame_time = dst::time::Instant::now();
1693                }
1694
1695                match read_result {
1696                    Ok(Some(Ok(Message::Binary(data)))) => {
1697                        log::trace!("Received message <binary> {} bytes", data.len());
1698                        last_data_time = dst::time::Instant::now();
1699
1700                        if !ConnectionMode::from_atomic(&connection_state).is_active()
1701                            || !read_fence.is_valid()
1702                        {
1703                            log::debug!(
1704                                "Dropping WebSocket message with {} bytes after session ended",
1705                                data.len()
1706                            );
1707                            break;
1708                        }
1709
1710                        if let Some(ref handler) = handler {
1711                            handler.handle(connection_epoch, Message::Binary(data));
1712                        }
1713                    }
1714                    Ok(Some(Ok(Message::Text(data)))) => {
1715                        log::trace!("Received text frame ({} bytes)", data.len());
1716                        last_data_time = dst::time::Instant::now();
1717
1718                        if !ConnectionMode::from_atomic(&connection_state).is_active()
1719                            || !read_fence.is_valid()
1720                        {
1721                            log::debug!(
1722                                "Dropping WebSocket message with {} bytes after session ended",
1723                                data.len()
1724                            );
1725                            break;
1726                        }
1727
1728                        if let Some(ref handler) = handler {
1729                            handler.handle(connection_epoch, Message::Text(data));
1730                        }
1731                    }
1732                    Ok(Some(Ok(Message::Ping(ping_data)))) => {
1733                        log::trace!("Received ping frame ({} bytes)", ping_data.len());
1734                        // Do not reset last_data_time: pings are keep-alive frames, not application
1735                        // data, so a peer that emits only pings must still trip the idle timeout.
1736                        // Checked here too: a ping flood faster than the check interval starves the timeout branch
1737
1738                        if let Some(ref handler) = ping_handler {
1739                            if !ConnectionMode::from_atomic(&connection_state).is_active()
1740                                || !read_fence.is_valid()
1741                            {
1742                                log::debug!(
1743                                    "Dropping WebSocket ping with {} bytes after session ended",
1744                                    ping_data.len()
1745                                );
1746                                break;
1747                            }
1748
1749                            handler.handle(connection_epoch, ping_data.to_vec());
1750                        }
1751
1752                        if idle_timeout_exceeded(last_data_time, idle_timeout) {
1753                            break;
1754                        }
1755                    }
1756                    Ok(Some(Ok(Message::Pong(_)))) => {
1757                        log::trace!("Received pong");
1758                        // Do not reset last_data_time: pongs are keep-alive replies (not data)
1759
1760                        if idle_timeout_exceeded(last_data_time, idle_timeout) {
1761                            break;
1762                        }
1763                    }
1764                    Ok(Some(Ok(Message::Close(frame)))) => {
1765                        Self::dispatch_close_frame(
1766                            &connection_state,
1767                            &read_fence,
1768                            close_handler.as_ref(),
1769                            state_sink.as_ref(),
1770                            frame,
1771                        );
1772                        break;
1773                    }
1774                    Ok(Some(Err(e))) => {
1775                        if is_connection_drop_transport_error(&e) {
1776                            log::warn!("Received connection error, terminating: {e}");
1777                        } else {
1778                            log::error!("Received transport error, terminating: {e}");
1779                        }
1780
1781                        break;
1782                    }
1783                    Ok(None) => {
1784                        log::log!(
1785                            read_termination_log_level(&connection_state),
1786                            "Connection closed by peer (no close frame), terminating"
1787                        );
1788                        break;
1789                    }
1790                    Err(_) => {
1791                        if heartbeat_timeout_exceeded(last_frame_time, heartbeat_timeout) {
1792                            break;
1793                        }
1794
1795                        if idle_timeout_exceeded(last_data_time, idle_timeout) {
1796                            break;
1797                        }
1798                    }
1799                }
1800            }
1801
1802            // Wake the controller immediately so it detects the dead read task
1803            state_notify.notify_one();
1804        })
1805    }
1806
1807    fn dispatch_close_frame(
1808        connection_state: &AtomicU8,
1809        read_fence: &ReadSessionFence,
1810        close_handler: Option<&CloseHandler>,
1811        state_sink: Option<&SocketStateSink>,
1812        frame: Option<CloseFrame>,
1813    ) {
1814        match &frame {
1815            Some(frame) => log::log!(
1816                read_termination_log_level(connection_state),
1817                "Received close frame, terminating: code={}, reason={REDACTED}",
1818                frame.code
1819            ),
1820            None => log::log!(
1821                read_termination_log_level(connection_state),
1822                "Received close frame with no code or reason, terminating"
1823            ),
1824        }
1825
1826        if !ConnectionMode::from_atomic(connection_state).is_active() || !read_fence.is_valid() {
1827            return;
1828        }
1829
1830        if let Some(handler) = close_handler
1831            && !handler(frame)
1832        {
1833            ConnectionMode::close_websocket_on_loss(connection_state, state_sink);
1834        }
1835    }
1836
1837    /// Queues a message for replay on the replacement connection.
1838    ///
1839    /// A keepalive belongs to the connection it was issued on: a Pong answers that
1840    /// connection's Ping, a heartbeat probes it, and a Close terminates it. None
1841    /// carries its meaning on a replacement connection, so those messages are dropped
1842    /// instead. Text keepalives reach the writer as [`WriterCommand::Heartbeat`] and
1843    /// never enter this buffer.
1844    fn buffer_for_replay(
1845        buffer: &mut VecDeque<(Message, tokio::sync::OwnedSemaphorePermit)>,
1846        msg: Message,
1847        permit: tokio::sync::OwnedSemaphorePermit,
1848    ) {
1849        if msg.is_control() {
1850            return;
1851        }
1852
1853        log::debug!(
1854            "Buffering message for replay (buffer size: {})",
1855            buffer.len() + 1
1856        );
1857
1858        buffer.push_back((msg, permit));
1859    }
1860
1861    /// Attempts to send all buffered messages after reconnection.
1862    ///
1863    /// Returns `true` if a send error occurred (caller should trigger reconnection).
1864    /// Messages remain in buffer if send fails, preserving them for the next reconnection attempt.
1865    async fn drain_reconnect_buffer(
1866        buffer: &mut VecDeque<(Message, tokio::sync::OwnedSemaphorePermit)>,
1867        writer: &mut MessageWriter,
1868        connection_state: &AtomicU8,
1869        auth_tracker: &Arc<OnceLock<AuthTracker>>,
1870        reconnect_buffer_waits_for_auth: &AtomicBool,
1871    ) -> bool {
1872        if buffer.is_empty() {
1873            return false;
1874        }
1875
1876        let initial_buffer_len = buffer.len();
1877        log::info!("Sending {initial_buffer_len} buffered messages after reconnection");
1878
1879        while !buffer.is_empty() {
1880            match Self::reconnect_buffer_action(
1881                reconnect_buffer_waits_for_auth,
1882                auth_tracker,
1883                connection_state,
1884            ) {
1885                ReconnectBufferAction::Drain => {}
1886                ReconnectBufferAction::Wait => return false,
1887                ReconnectBufferAction::Discard => {
1888                    log::warn!(
1889                        "Discarding {} buffered messages after authentication failed",
1890                        buffer.len()
1891                    );
1892                    buffer.clear();
1893                    return false;
1894                }
1895            }
1896
1897            // Clone message before attempting send (to keep in buffer if send fails)
1898            let msg_to_send = buffer
1899                .front()
1900                .expect("reconnect buffer should not be empty")
1901                .0
1902                .clone();
1903
1904            if let Err(e) = writer.send(msg_to_send).await {
1905                if is_connection_drop_transport_error(&e) {
1906                    log::warn!(
1907                        "Failed to send buffered message after reconnection: {e}, {} messages remain in buffer",
1908                        buffer.len()
1909                    );
1910                } else {
1911                    log::error!(
1912                        "Failed to send buffered message after reconnection: {e}, {} messages remain in buffer",
1913                        buffer.len()
1914                    );
1915                }
1916
1917                return true;
1918            }
1919
1920            // Only remove from buffer after successful send
1921            buffer.pop_front();
1922        }
1923
1924        if buffer.is_empty() {
1925            log::info!("Successfully sent all {initial_buffer_len} buffered messages");
1926        }
1927
1928        false
1929    }
1930
1931    fn can_drain_reconnect_buffer(
1932        reconnect_buffer_waits_for_auth: &AtomicBool,
1933        auth_tracker: &Arc<OnceLock<AuthTracker>>,
1934    ) -> ReconnectBufferAction {
1935        if !reconnect_buffer_waits_for_auth.load(Ordering::Acquire) {
1936            return ReconnectBufferAction::Drain;
1937        }
1938
1939        match auth_tracker.get().map(AuthTracker::auth_state) {
1940            Some(AuthState::Authenticated) => ReconnectBufferAction::Drain,
1941            Some(AuthState::Failed) => ReconnectBufferAction::Discard,
1942            Some(AuthState::Unauthenticated) | None => ReconnectBufferAction::Wait,
1943        }
1944    }
1945
1946    fn reconnect_buffer_action(
1947        reconnect_buffer_waits_for_auth: &AtomicBool,
1948        auth_tracker: &Arc<OnceLock<AuthTracker>>,
1949        connection_state: &AtomicU8,
1950    ) -> ReconnectBufferAction {
1951        let action =
1952            Self::can_drain_reconnect_buffer(reconnect_buffer_waits_for_auth, auth_tracker);
1953
1954        // Keep mode as the final admission check so an accepted reconnect stops the next send.
1955        if action == ReconnectBufferAction::Drain
1956            && !ConnectionMode::from_atomic(connection_state).is_active()
1957        {
1958            ReconnectBufferAction::Wait
1959        } else {
1960            action
1961        }
1962    }
1963
1964    #[expect(
1965        clippy::too_many_arguments,
1966        reason = "writer task owns the transport and shared lifecycle coordination state"
1967    )]
1968    fn spawn_write_task(
1969        connection_state: Arc<AtomicU8>,
1970        controller_notify: Arc<tokio::sync::Notify>,
1971        reconnect_published: Arc<AtomicBool>,
1972        writer: MessageWriter,
1973        mut writer_rx: WriterReceiver<WriterCommand>,
1974        connection_epoch: Arc<AtomicU64>,
1975        auth_tracker: Arc<OnceLock<AuthTracker>>,
1976        reconnect_buffer_waits_for_auth: Arc<AtomicBool>,
1977        state_sink: Option<SocketStateSink>,
1978    ) -> dst::task::JoinHandle<()> {
1979        log_task_started("write");
1980
1981        // Interval between checking the connection mode
1982        let check_interval = Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS);
1983
1984        dst::task::spawn(async move {
1985            let mut active_writer = writer;
1986            // Buffer for messages received during reconnection
1987            // VecDeque for efficient pop_front() operations
1988            let mut reconnect_buffer: VecDeque<(Message, tokio::sync::OwnedSemaphorePermit)> =
1989                VecDeque::new();
1990
1991            loop {
1992                let mode = ConnectionMode::from_atomic(&connection_state);
1993
1994                match mode {
1995                    ConnectionMode::Disconnect => {
1996                        // Log any buffered messages that will be lost
1997                        if !reconnect_buffer.is_empty() {
1998                            log::warn!(
1999                                "Discarding {} buffered messages due to disconnect",
2000                                reconnect_buffer.len()
2001                            );
2002                            reconnect_buffer.clear();
2003                        }
2004
2005                        // Attempt to close the writer gracefully before exiting,
2006                        // we ignore any error as the writer may already be closed.
2007                        _ = dst::time::timeout(
2008                            Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS),
2009                            active_writer.close(),
2010                        )
2011                        .await;
2012                        break;
2013                    }
2014                    ConnectionMode::Closed => {
2015                        // Log any buffered messages that will be lost
2016                        if !reconnect_buffer.is_empty() {
2017                            log::warn!(
2018                                "Discarding {} buffered messages due to closed connection",
2019                                reconnect_buffer.len()
2020                            );
2021                            reconnect_buffer.clear();
2022                        }
2023
2024                        break;
2025                    }
2026                    _ => {}
2027                }
2028
2029                if mode.is_active() && !reconnect_buffer.is_empty() {
2030                    match Self::reconnect_buffer_action(
2031                        reconnect_buffer_waits_for_auth.as_ref(),
2032                        &auth_tracker,
2033                        &connection_state,
2034                    ) {
2035                        ReconnectBufferAction::Drain => {
2036                            let drain_result = dst::time::timeout(
2037                                Duration::from_secs(WRITE_TIMEOUT_SECS),
2038                                Self::drain_reconnect_buffer(
2039                                    &mut reconnect_buffer,
2040                                    &mut active_writer,
2041                                    &connection_state,
2042                                    &auth_tracker,
2043                                    reconnect_buffer_waits_for_auth.as_ref(),
2044                                ),
2045                            )
2046                            .await;
2047
2048                            let send_error = drain_result.unwrap_or_else(|_| {
2049                                log::warn!(
2050                                    "Timed out draining reconnect buffer after {WRITE_TIMEOUT_SECS}s, {} messages remain",
2051                                    reconnect_buffer.len()
2052                                );
2053                                true
2054                            });
2055
2056                            // CAS: a disconnect landing mid-drain must not be overwritten
2057                            if send_error {
2058                                _ = request_websocket_reconnect(
2059                                    &connection_state,
2060                                    &reconnect_published,
2061                                    state_sink.as_ref(),
2062                                    &auth_tracker,
2063                                    &controller_notify,
2064                                    || {},
2065                                );
2066                            }
2067
2068                            continue;
2069                        }
2070                        ReconnectBufferAction::Discard => {
2071                            log::warn!(
2072                                "Discarding {} buffered messages after authentication failed",
2073                                reconnect_buffer.len()
2074                            );
2075                            reconnect_buffer.clear();
2076                            continue;
2077                        }
2078                        ReconnectBufferAction::Wait => {}
2079                    }
2080                }
2081
2082                match dst::time::timeout(check_interval, writer_rx.recv()).await {
2083                    Ok(Some((msg, permit))) => {
2084                        // Re-check connection mode after receiving a message
2085                        let mode = ConnectionMode::from_atomic(&connection_state);
2086
2087                        if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
2088                            break;
2089                        }
2090
2091                        match msg {
2092                            WriterCommand::Update(new_writer, tx) => {
2093                                log::debug!("Received new writer");
2094
2095                                // Delay before closing connection
2096                                dst::time::sleep(Duration::from_millis(100)).await;
2097
2098                                // Attempt to close the writer gracefully on update,
2099                                // we ignore any error as the writer may already be closed.
2100                                _ = dst::time::timeout(
2101                                    Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS),
2102                                    active_writer.close(),
2103                                )
2104                                .await;
2105
2106                                active_writer = new_writer;
2107                                let epoch = connection_epoch.fetch_add(1, Ordering::AcqRel) + 1;
2108                                log::debug!("Updated writer: epoch={epoch}");
2109
2110                                if let Err(e) = tx.send(epoch) {
2111                                    log::error!(
2112                                        "Failed to report writer update to controller: {e:?}"
2113                                    );
2114                                }
2115                            }
2116                            WriterCommand::Send(msg)
2117                                if mode.is_reconnect()
2118                                    || (!reconnect_buffer.is_empty() && !msg.is_control()) =>
2119                            {
2120                                Self::buffer_for_replay(&mut reconnect_buffer, msg, permit);
2121                            }
2122                            WriterCommand::Heartbeat(_)
2123                            | WriterCommand::SendPongOnConnection { .. }
2124                                if mode.is_reconnect() => {}
2125                            WriterCommand::SendOnConnection { response_tx, .. }
2126                                if mode.is_reconnect() =>
2127                            {
2128                                _ = response_tx.send(Err(SendError::ConnectionChanged));
2129                            }
2130                            WriterCommand::SendPongOnConnection {
2131                                data,
2132                                connection_epoch: expected_epoch,
2133                            } => {
2134                                let epoch = connection_epoch.load(Ordering::Acquire);
2135                                if epoch != expected_epoch {
2136                                    continue;
2137                                }
2138
2139                                let send_result = dst::time::timeout(
2140                                    Duration::from_secs(WRITE_TIMEOUT_SECS),
2141                                    active_writer.send(Message::Pong(data.into())),
2142                                )
2143                                .await;
2144
2145                                let send_failed = match send_result {
2146                                    Ok(Ok(())) => false,
2147                                    Ok(Err(e)) => {
2148                                        if is_connection_drop_transport_error(&e) {
2149                                            log::warn!("Failed to send pong: {e}");
2150                                        } else {
2151                                            log::error!("Failed to send pong: {e}");
2152                                        }
2153
2154                                        true
2155                                    }
2156                                    Err(_) => {
2157                                        log::warn!(
2158                                            "Timed out sending pong after {WRITE_TIMEOUT_SECS}s"
2159                                        );
2160                                        true
2161                                    }
2162                                };
2163
2164                                if send_failed
2165                                    && request_websocket_reconnect(
2166                                        &connection_state,
2167                                        &reconnect_published,
2168                                        state_sink.as_ref(),
2169                                        &auth_tracker,
2170                                        &controller_notify,
2171                                        || {},
2172                                    ) == ReconnectRequestOutcome::Accepted
2173                                {
2174                                    log::warn!("Writer triggering reconnect");
2175                                }
2176                            }
2177                            WriterCommand::SendOnConnection {
2178                                message,
2179                                connection_epoch: expected_epoch,
2180                                response_tx,
2181                            } => {
2182                                let epoch = connection_epoch.load(Ordering::Acquire);
2183                                if epoch != expected_epoch {
2184                                    _ = response_tx.send(Err(SendError::ConnectionChanged));
2185                                    continue;
2186                                }
2187
2188                                let send_result = dst::time::timeout(
2189                                    Duration::from_secs(WRITE_TIMEOUT_SECS),
2190                                    active_writer.send(message),
2191                                )
2192                                .await;
2193
2194                                // An ownership-bound message is never replayed, so an expired
2195                                // deadline reports failure to the caller instead of buffering
2196                                // the message as the ordinary send path does.
2197                                let result = match send_result {
2198                                    Ok(Ok(())) => Ok(()),
2199                                    Ok(Err(e)) => {
2200                                        if is_connection_drop_transport_error(&e) {
2201                                            log::warn!("Failed to send message: {e}");
2202                                        } else {
2203                                            log::error!("Failed to send message: {e}");
2204                                        }
2205
2206                                        Err(SendError::BrokenPipe(e.to_string()))
2207                                    }
2208                                    Err(_) => {
2209                                        log::warn!(
2210                                            "Timed out sending message after {WRITE_TIMEOUT_SECS}s"
2211                                        );
2212
2213                                        Err(SendError::WriteTimeout)
2214                                    }
2215                                };
2216
2217                                let send_failed = result.is_err();
2218                                _ = response_tx.send(result);
2219
2220                                if send_failed
2221                                    && request_websocket_reconnect(
2222                                        &connection_state,
2223                                        &reconnect_published,
2224                                        state_sink.as_ref(),
2225                                        &auth_tracker,
2226                                        &controller_notify,
2227                                        || {},
2228                                    ) == ReconnectRequestOutcome::Accepted
2229                                {
2230                                    log::warn!("Writer triggering reconnect");
2231                                }
2232                            }
2233                            WriterCommand::Send(msg) => {
2234                                let send_failed =
2235                                    Self::write_outbound(&mut active_writer, msg.clone()).await;
2236                                if send_failed {
2237                                    Self::buffer_for_replay(&mut reconnect_buffer, msg, permit);
2238
2239                                    // CAS: a disconnect landing mid-send must not be overwritten
2240                                    if request_websocket_reconnect(
2241                                        &connection_state,
2242                                        &reconnect_published,
2243                                        state_sink.as_ref(),
2244                                        &auth_tracker,
2245                                        &controller_notify,
2246                                        || {},
2247                                    ) == ReconnectRequestOutcome::Accepted
2248                                    {
2249                                        log::warn!("Writer triggering reconnect");
2250                                    }
2251                                }
2252                            }
2253                            WriterCommand::Heartbeat(msg) => {
2254                                let send_failed =
2255                                    Self::write_outbound(&mut active_writer, msg).await;
2256
2257                                if send_failed
2258                                    && request_websocket_reconnect(
2259                                        &connection_state,
2260                                        &reconnect_published,
2261                                        state_sink.as_ref(),
2262                                        &auth_tracker,
2263                                        &controller_notify,
2264                                        || {},
2265                                    ) == ReconnectRequestOutcome::Accepted
2266                                {
2267                                    log::warn!("Writer triggering reconnect");
2268                                }
2269                            }
2270                        }
2271                    }
2272                    Ok(None) => {
2273                        // Channel closed - writer task should terminate
2274                        log::debug!("Writer channel closed, terminating writer task");
2275                        break;
2276                    }
2277                    Err(_) => {
2278                        // Timeout - just continue the loop
2279                    }
2280                }
2281            }
2282
2283            // Attempt to close the writer gracefully before exiting,
2284            // we ignore any error as the writer may already be closed.
2285            _ = dst::time::timeout(
2286                Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS),
2287                active_writer.close(),
2288            )
2289            .await;
2290
2291            log_task_stopped("write");
2292        })
2293    }
2294
2295    async fn write_outbound(writer: &mut MessageWriter, msg: Message) -> bool {
2296        let send_result =
2297            dst::time::timeout(Duration::from_secs(WRITE_TIMEOUT_SECS), writer.send(msg)).await;
2298
2299        match send_result {
2300            Ok(Ok(())) => false,
2301            Ok(Err(e)) => {
2302                if is_connection_drop_transport_error(&e) {
2303                    log::warn!("Failed to send message: {e}");
2304                } else {
2305                    log::error!("Failed to send message: {e}");
2306                }
2307
2308                true
2309            }
2310            Err(_) => {
2311                log::warn!("Timed out sending message after {WRITE_TIMEOUT_SECS}s");
2312                true
2313            }
2314        }
2315    }
2316
2317    fn spawn_heartbeat_task(
2318        connection_state: Arc<AtomicU8>,
2319        heartbeat_secs: u64,
2320        message: Option<String>,
2321        writer_tx: WriterSender<WriterCommand>,
2322    ) -> dst::task::JoinHandle<()> {
2323        log_task_started("heartbeat");
2324
2325        dst::task::spawn(async move {
2326            let interval = Duration::from_secs(heartbeat_secs);
2327            loop {
2328                dst::time::sleep(interval).await;
2329
2330                match ConnectionMode::from_u8(connection_state.load(Ordering::SeqCst)) {
2331                    ConnectionMode::Active => {
2332                        let msg = match &message {
2333                            Some(text) => {
2334                                WriterCommand::Heartbeat(Message::Text(text.clone().into()))
2335                            }
2336                            None => WriterCommand::Heartbeat(Message::Ping(vec![].into())),
2337                        };
2338
2339                        match writer_tx.send_control(msg) {
2340                            Ok(()) => log::trace!("Sent heartbeat to writer task"),
2341                            Err(e) => {
2342                                log::warn!("Failed to send heartbeat to writer task: {e}");
2343                            }
2344                        }
2345                    }
2346                    ConnectionMode::Reconnect => {}
2347                    ConnectionMode::Disconnect | ConnectionMode::Closed => break,
2348                }
2349            }
2350
2351            log_task_stopped("heartbeat");
2352        })
2353    }
2354}
2355
2356fn heartbeat_timeout_exceeded(
2357    last_frame_time: dst::time::Instant,
2358    timeout: Option<Duration>,
2359) -> bool {
2360    if let Some(timeout) = timeout {
2361        let elapsed = last_frame_time.elapsed();
2362        if elapsed >= timeout {
2363            log::warn!(
2364                "Heartbeat timeout: no frame received for {:.1}s",
2365                elapsed.as_secs_f64()
2366            );
2367            return true;
2368        }
2369    }
2370
2371    false
2372}
2373
2374fn idle_timeout_exceeded(
2375    last_data_time: dst::time::Instant,
2376    idle_timeout: Option<Duration>,
2377) -> bool {
2378    if let Some(timeout) = idle_timeout {
2379        let idle_duration = last_data_time.elapsed();
2380        if idle_duration >= timeout {
2381            log::warn!(
2382                "Read idle timeout: no data received for {:.1}s",
2383                idle_duration.as_secs_f64()
2384            );
2385            return true;
2386        }
2387    }
2388
2389    false
2390}
2391
2392impl Drop for WebSocketClientInner {
2393    fn drop(&mut self) {
2394        if let Some(read_fence) = self.read_fence.take() {
2395            read_fence.invalidate();
2396        }
2397
2398        if let Some(ref read_task) = self.read_task.take()
2399            && !read_task.is_finished()
2400        {
2401            read_task.abort();
2402            log_task_aborted("read");
2403        }
2404
2405        if !self.write_task.is_finished() {
2406            self.write_task.abort();
2407            log_task_aborted("write");
2408        }
2409
2410        if let Some(ref handle) = self.heartbeat_task.take()
2411            && !handle.is_finished()
2412        {
2413            handle.abort();
2414            log_task_aborted("heartbeat");
2415        }
2416    }
2417}
2418
2419#[expect(
2420    clippy::missing_fields_in_debug,
2421    reason = "handler closures and internal task handles are intentionally omitted"
2422)]
2423impl Debug for WebSocketClientInner {
2424    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2425        f.debug_struct(stringify!(WebSocketClientInner))
2426            .field("config", &self.config)
2427            .field(
2428                "connection_mode",
2429                &ConnectionMode::from_atomic(&self.connection_mode),
2430            )
2431            .field("connect_timeout", &self.connect_timeout)
2432            .field("is_stream_mode", &self.handler.is_none())
2433            .finish()
2434    }
2435}
2436
2437#[derive(Clone)]
2438enum IncomingHandler {
2439    Message(MessageHandler),
2440    Epoch(EpochMessageHandler),
2441}
2442
2443impl IncomingHandler {
2444    fn handle(&self, connection_epoch: u64, message: Message) {
2445        match self {
2446            Self::Message(handler) => handler(message),
2447            Self::Epoch(handler) => handler(connection_epoch, message),
2448        }
2449    }
2450}
2451
2452#[derive(Clone)]
2453enum IncomingPingHandler {
2454    Ping(PingHandler),
2455    Epoch(EpochPingHandler),
2456}
2457
2458impl IncomingPingHandler {
2459    fn handle(&self, connection_epoch: u64, data: Vec<u8>) {
2460        match self {
2461            Self::Ping(handler) => handler(data),
2462            Self::Epoch(handler) => handler(connection_epoch, data),
2463        }
2464    }
2465}
2466
2467#[derive(Clone, Copy, Debug, Eq, PartialEq)]
2468enum ReconnectBufferAction {
2469    Drain,
2470    Wait,
2471    Discard,
2472}
2473
2474/// A WebSocket client with rate limiting, heartbeats, and automatic reconnection in handler mode.
2475///
2476/// Handler mode owns the reader and writer tasks, buffers sends during reconnection, and replays
2477/// them against the replacement connection. Stream mode returns the reader to the caller and does
2478/// not reconnect automatically. See [`crate::websocket`] for connection ownership, replay, and
2479/// epoch guarantees.
2480pub struct WebSocketClient {
2481    pub(crate) controller_task: dst::task::JoinHandle<()>,
2482    pub(crate) connection_mode: Arc<AtomicU8>,
2483    pub(crate) connection_epoch: Arc<AtomicU64>,
2484    pub(crate) state_notify: Arc<tokio::sync::Notify>,
2485    pub(crate) connect_timeout: Duration,
2486    pub(crate) rate_limiter: Arc<RateLimiter<Ustr, MonotonicClock>>,
2487    pub(crate) writer_tx: WriterSender<WriterCommand>,
2488    auth_tracker: Arc<OnceLock<AuthTracker>>,
2489    reconnect_buffer_waits_for_auth: Arc<AtomicBool>,
2490    reconnect_headers: ReconnectHeaders,
2491    state_sink: Option<SocketStateSink>,
2492    controller_lifecycle: Arc<ControllerLifecycle>,
2493    controller_notify: Arc<tokio::sync::Notify>,
2494    reconnect_published: Arc<AtomicBool>,
2495    reconnect_supported: bool,
2496}
2497
2498/// Shared headers used by future automatic WebSocket reconnects.
2499///
2500/// Updating these headers does not affect the active connection or trigger a reconnect.
2501#[derive(Clone)]
2502pub struct ReconnectHeaders {
2503    inner: Arc<RwLock<Vec<(String, String)>>>,
2504}
2505
2506impl ReconnectHeaders {
2507    fn new(headers: Vec<(String, String)>) -> Self {
2508        Self {
2509            inner: Arc::new(RwLock::new(headers)),
2510        }
2511    }
2512
2513    /// Replaces a header used by future automatic reconnect attempts.
2514    ///
2515    /// # Errors
2516    ///
2517    /// Returns an error if the header name or value is invalid.
2518    pub fn update(&self, name: &str, value: &str) -> Result<(), TransportError> {
2519        let name = HeaderName::from_bytes(name.as_bytes()).map_err(|e| {
2520            TransportError::Io(std::io::Error::new(
2521                std::io::ErrorKind::InvalidInput,
2522                format!("Invalid WebSocket reconnect header name: {e}"),
2523            ))
2524        })?;
2525
2526        HeaderValue::from_str(value).map_err(|e| {
2527            TransportError::Io(std::io::Error::new(
2528                std::io::ErrorKind::InvalidInput,
2529                format!("Invalid WebSocket reconnect header value: {e}"),
2530            ))
2531        })?;
2532
2533        let name = name.as_str();
2534        let mut headers = self.inner.write();
2535        headers.retain(|(existing, _)| !existing.eq_ignore_ascii_case(name));
2536        headers.push((name.to_string(), value.to_string()));
2537        Ok(())
2538    }
2539
2540    fn snapshot(&self) -> Vec<(String, String)> {
2541        self.inner.read().clone()
2542    }
2543}
2544
2545impl Debug for ReconnectHeaders {
2546    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2547        f.debug_struct(stringify!(ReconnectHeaders))
2548            .finish_non_exhaustive()
2549    }
2550}
2551
2552/// Cloneable controller handle for requesting one transport reconnect.
2553#[derive(Clone)]
2554pub struct WebSocketReconnectHandle {
2555    connection_mode: Arc<AtomicU8>,
2556    auth_tracker: Arc<OnceLock<AuthTracker>>,
2557    state_sink: Option<SocketStateSink>,
2558    controller_lifecycle: Arc<ControllerLifecycle>,
2559    controller_notify: Arc<tokio::sync::Notify>,
2560    reconnect_published: Arc<AtomicBool>,
2561    supported: bool,
2562}
2563
2564impl Debug for WebSocketReconnectHandle {
2565    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2566        f.debug_struct(stringify!(WebSocketReconnectHandle))
2567            .field(
2568                "connection_mode",
2569                &ConnectionMode::from_atomic(&self.connection_mode),
2570            )
2571            .field("supported", &self.supported)
2572            .finish_non_exhaustive()
2573    }
2574}
2575
2576impl WebSocketReconnectHandle {
2577    /// Requests that the controller replace the active transport.
2578    ///
2579    /// An accepted request invalidates registered authentication state, wakes the controller, and
2580    /// publishes the loss before the replacement can become active. Rejected requests leave the
2581    /// transport and authentication state unchanged. Handler-mode handles return
2582    /// [`ReconnectRequestOutcome::Closed`] after the client drops; stream-mode handles always
2583    /// return [`ReconnectRequestOutcome::Unsupported`].
2584    #[must_use]
2585    pub fn request_reconnect(&self) -> ReconnectRequestOutcome {
2586        if !self.supported {
2587            return ReconnectRequestOutcome::Unsupported;
2588        }
2589
2590        let Some(request) = self.controller_lifecycle.enter_request() else {
2591            return ReconnectRequestOutcome::Closed;
2592        };
2593
2594        let mut request = Some(request);
2595        request_websocket_reconnect(
2596            &self.connection_mode,
2597            &self.reconnect_published,
2598            self.state_sink.as_ref(),
2599            &self.auth_tracker,
2600            &self.controller_notify,
2601            || drop(request.take()),
2602        )
2603    }
2604}
2605
2606struct ReconnectPublication<'a> {
2607    published: &'a AtomicBool,
2608    controller_notify: &'a tokio::sync::Notify,
2609}
2610
2611impl Drop for ReconnectPublication<'_> {
2612    fn drop(&mut self) {
2613        self.published.store(true, Ordering::SeqCst);
2614        self.controller_notify.notify_one();
2615    }
2616}
2617
2618fn request_websocket_reconnect<F>(
2619    connection_mode: &AtomicU8,
2620    reconnect_published: &AtomicBool,
2621    state_sink: Option<&SocketStateSink>,
2622    auth_tracker: &OnceLock<AuthTracker>,
2623    controller_notify: &tokio::sync::Notify,
2624    on_handoff: F,
2625) -> ReconnectRequestOutcome
2626where
2627    F: FnOnce(),
2628{
2629    if reconnect_published
2630        .compare_exchange(true, false, Ordering::SeqCst, Ordering::SeqCst)
2631        .is_err()
2632    {
2633        return match ConnectionMode::from_atomic(connection_mode) {
2634            ConnectionMode::Active | ConnectionMode::Reconnect => {
2635                ReconnectRequestOutcome::AlreadyReconnecting
2636            }
2637            ConnectionMode::Disconnect => ReconnectRequestOutcome::Disconnected,
2638            ConnectionMode::Closed => ReconnectRequestOutcome::Closed,
2639        };
2640    }
2641
2642    let outcome = ConnectionMode::request_reconnect_outcome(connection_mode);
2643    if outcome != ReconnectRequestOutcome::Accepted {
2644        reconnect_published.store(true, Ordering::SeqCst);
2645        return outcome;
2646    }
2647
2648    let _publication = ReconnectPublication {
2649        published: reconnect_published,
2650        controller_notify,
2651    };
2652
2653    if let Some(tracker) = auth_tracker.get() {
2654        tracker.invalidate();
2655    }
2656
2657    controller_notify.notify_one();
2658    on_handoff();
2659
2660    if let Some(sink) = state_sink {
2661        sink.publish_websocket(SocketState::Disconnected);
2662    }
2663
2664    ReconnectRequestOutcome::Accepted
2665}
2666
2667#[cfg(test)]
2668mod reconnect_request_tests {
2669    use std::sync::{Arc, OnceLock, atomic::AtomicU8};
2670
2671    use rstest::rstest;
2672
2673    use super::*;
2674
2675    fn handle(
2676        mode: ConnectionMode,
2677        supported: bool,
2678    ) -> (
2679        WebSocketReconnectHandle,
2680        AuthTracker,
2681        Arc<tokio::sync::Notify>,
2682    ) {
2683        let tracker = AuthTracker::new();
2684        let _receiver = tracker.begin();
2685        tracker.succeed();
2686        let auth_tracker = Arc::new(OnceLock::new());
2687        auth_tracker
2688            .set(tracker.clone())
2689            .expect("auth tracker should be unset");
2690        let notify = Arc::new(tokio::sync::Notify::new());
2691
2692        let handle = WebSocketReconnectHandle {
2693            connection_mode: Arc::new(AtomicU8::new(mode.as_u8())),
2694            auth_tracker,
2695            state_sink: None,
2696            controller_lifecycle: Arc::new(ControllerLifecycle::new()),
2697            controller_notify: Arc::clone(&notify),
2698            reconnect_published: Arc::new(AtomicBool::new(true)),
2699            supported,
2700        };
2701
2702        (handle, tracker, notify)
2703    }
2704
2705    #[rstest]
2706    #[tokio::test]
2707    async fn accepted_request_invalidates_auth_and_wakes_controller_once() {
2708        let (handle, tracker, notify) = handle(ConnectionMode::Active, true);
2709
2710        assert_eq!(
2711            handle.request_reconnect(),
2712            ReconnectRequestOutcome::Accepted
2713        );
2714        assert!(!tracker.is_authenticated());
2715        tokio::time::timeout(Duration::from_millis(10), notify.notified())
2716            .await
2717            .expect("accepted request should notify controller");
2718
2719        let _receiver = tracker.begin();
2720        tracker.succeed();
2721        assert_eq!(
2722            handle.request_reconnect(),
2723            ReconnectRequestOutcome::AlreadyReconnecting
2724        );
2725        assert!(tracker.is_authenticated());
2726        assert!(
2727            tokio::time::timeout(Duration::from_millis(10), notify.notified())
2728                .await
2729                .is_err(),
2730            "duplicate request should not notify controller",
2731        );
2732
2733        handle
2734            .connection_mode
2735            .store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
2736        assert_eq!(
2737            handle.request_reconnect(),
2738            ReconnectRequestOutcome::Accepted
2739        );
2740        assert!(!tracker.is_authenticated());
2741        tokio::time::timeout(Duration::from_millis(10), notify.notified())
2742            .await
2743            .expect("a later reconnect cycle should notify the controller");
2744    }
2745
2746    #[rstest]
2747    #[case(
2748        ConnectionMode::Disconnect,
2749        true,
2750        ReconnectRequestOutcome::Disconnected
2751    )]
2752    #[case(
2753        ConnectionMode::Reconnect,
2754        true,
2755        ReconnectRequestOutcome::AlreadyReconnecting
2756    )]
2757    #[case(ConnectionMode::Closed, true, ReconnectRequestOutcome::Closed)]
2758    #[case(ConnectionMode::Active, false, ReconnectRequestOutcome::Unsupported)]
2759    #[tokio::test]
2760    async fn rejected_request_preserves_auth_and_does_not_wake_controller(
2761        #[case] mode: ConnectionMode,
2762        #[case] supported: bool,
2763        #[case] expected: ReconnectRequestOutcome,
2764    ) {
2765        let (handle, tracker, notify) = handle(mode, supported);
2766
2767        assert_eq!(handle.request_reconnect(), expected);
2768        assert!(tracker.is_authenticated());
2769        assert!(
2770            tokio::time::timeout(Duration::from_millis(10), notify.notified())
2771                .await
2772                .is_err(),
2773            "rejected request should not notify controller",
2774        );
2775    }
2776
2777    #[rstest]
2778    #[tokio::test]
2779    async fn reconnect_loss_callback_rejects_nested_request() {
2780        let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
2781        let controller_notify = Arc::new(tokio::sync::Notify::new());
2782        let handle_slot = Arc::new(OnceLock::<WebSocketReconnectHandle>::new());
2783        let handle_slot_callback = Arc::clone(&handle_slot);
2784        let nested_outcomes = Arc::new(parking_lot::Mutex::new(Vec::new()));
2785        let nested_outcomes_callback = Arc::clone(&nested_outcomes);
2786        let states = Arc::new(parking_lot::Mutex::new(Vec::new()));
2787        let states_callback = Arc::clone(&states);
2788        let sink = SocketStateSink::new(move |state| {
2789            states_callback.lock().push(state);
2790            nested_outcomes_callback
2791                .lock()
2792                .push(handle_slot_callback.get().unwrap().request_reconnect());
2793        });
2794
2795        let handle = WebSocketReconnectHandle {
2796            connection_mode,
2797            auth_tracker: Arc::new(OnceLock::new()),
2798            state_sink: Some(sink),
2799            controller_lifecycle: Arc::new(ControllerLifecycle::new()),
2800            controller_notify,
2801            reconnect_published: Arc::new(AtomicBool::new(true)),
2802            supported: true,
2803        };
2804
2805        handle_slot.set(handle.clone()).unwrap();
2806        let (result_tx, result_rx) = tokio::sync::oneshot::channel();
2807        std::thread::spawn(move || {
2808            _ = result_tx.send(handle.request_reconnect());
2809        });
2810
2811        assert_eq!(
2812            tokio::time::timeout(Duration::from_secs(1), result_rx)
2813                .await
2814                .expect("reentrant reconnect callback deadlocked")
2815                .unwrap(),
2816            ReconnectRequestOutcome::Accepted
2817        );
2818        assert_eq!(
2819            *nested_outcomes.lock(),
2820            vec![ReconnectRequestOutcome::AlreadyReconnecting]
2821        );
2822        assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
2823    }
2824
2825    #[rstest]
2826    fn closed_stream_handle_remains_unsupported() {
2827        let (handle, tracker, _notify) = handle(ConnectionMode::Closed, false);
2828        handle.controller_lifecycle.close_and_abort();
2829
2830        assert_eq!(
2831            handle.request_reconnect(),
2832            ReconnectRequestOutcome::Unsupported
2833        );
2834        assert!(tracker.is_authenticated());
2835    }
2836}
2837
2838impl Debug for WebSocketClient {
2839    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
2840        f.debug_struct(stringify!(WebSocketClient)).finish()
2841    }
2842}
2843
2844#[bon::bon]
2845impl WebSocketClient {
2846    /// Returns a builder for a websocket client in **stream mode**.
2847    ///
2848    /// Calling `connect` returns a stream that the caller owns and reads from directly. Automatic
2849    /// reconnection is **disabled** because the reader cannot be replaced internally. On
2850    /// disconnection, the client transitions to CLOSED state and the caller must manually create a
2851    /// new connection.
2852    ///
2853    /// Use stream mode when you need custom reconnection logic, direct control over message
2854    /// reading, or fine-grained backpressure handling.
2855    ///
2856    /// `default_quota` and `keyed_quotas` limit outgoing messages. `state_sink` reports transport
2857    /// availability changes.
2858    ///
2859    /// See [`WebSocketConfig`] documentation for comparison with handler mode.
2860    ///
2861    /// # Errors
2862    ///
2863    /// Returns an error if the connection cannot be established.
2864    #[builder(
2865        builder_type = WebSocketClientStreamBuilder,
2866        finish_fn = connect
2867    )]
2868    pub async fn stream_builder(
2869        config: WebSocketConfig,
2870        #[builder(default)] keyed_quotas: Vec<(String, Quota)>,
2871        default_quota: Option<Quota>,
2872        state_sink: Option<SocketStateSink>,
2873    ) -> Result<(MessageReader, Self), TransportError> {
2874        install_cryptographic_provider();
2875        validate_honored_inbound_limits(&config)?;
2876
2877        // Create a single connection and split it, respecting configured headers.
2878        // The connection attempt bound is a fixed default: stream mode documents reconnect_* fields as ignored
2879        let connect_timeout = Duration::from_secs(10);
2880
2881        let (writer, reader) = dst::time::timeout(
2882            connect_timeout,
2883            Box::pin(WebSocketClientInner::connect_with_server(
2884                &config.url,
2885                config.headers.clone(),
2886                config.backend,
2887                config.proxy_url.as_deref(),
2888                config.max_message_size_bytes,
2889                config.max_frame_size_bytes,
2890            )),
2891        )
2892        .await
2893        .map_err(|_| {
2894            TransportError::Io(std::io::Error::new(
2895                std::io::ErrorKind::TimedOut,
2896                format!(
2897                    "connection timed out after {}s",
2898                    connect_timeout.as_secs_f64()
2899                ),
2900            ))
2901        })??;
2902
2903        // Create inner without connecting (we'll provide the writer)
2904        let inner =
2905            WebSocketClientInner::new_with_writer_and_state_sink(config, writer, state_sink)?;
2906
2907        let connection_mode = Arc::clone(&inner.connection_mode);
2908        let connection_epoch = Arc::clone(&inner.connection_epoch);
2909        let state_notify = Arc::clone(&inner.state_notify);
2910        let controller_notify = Arc::clone(&inner.controller_notify);
2911        let reconnect_published = Arc::clone(&inner.reconnect_published);
2912        let connect_timeout = inner.connect_timeout;
2913        let auth_tracker = Arc::clone(&inner.auth_tracker);
2914        let reconnect_buffer_waits_for_auth = Arc::clone(&inner.reconnect_buffer_waits_for_auth);
2915        let reconnect_headers = inner.reconnect_headers.clone();
2916        let state_sink = inner.state_sink.clone();
2917        let rate_limiter = Self::resolve_rate_limiter(default_quota, keyed_quotas, None)?;
2918        let writer_tx = inner.writer_tx.clone();
2919        let controller_lifecycle = Arc::new(ControllerLifecycle::new());
2920
2921        let controller_task = Self::spawn_controller_task(
2922            inner,
2923            Arc::clone(&connection_mode),
2924            Arc::clone(&state_notify),
2925            Arc::clone(&auth_tracker),
2926            Arc::clone(&controller_lifecycle),
2927            Arc::clone(&controller_notify),
2928            Arc::clone(&reconnect_published),
2929        );
2930        let abort_handle = controller_task.abort_handle();
2931        controller_lifecycle.set_abort(move || abort_handle.abort());
2932
2933        Ok((
2934            reader,
2935            Self {
2936                controller_task,
2937                connection_mode,
2938                connection_epoch,
2939                state_notify,
2940                connect_timeout,
2941                rate_limiter,
2942                writer_tx,
2943                auth_tracker,
2944                reconnect_buffer_waits_for_auth,
2945                reconnect_headers,
2946                state_sink,
2947                controller_lifecycle,
2948                controller_notify,
2949                reconnect_published,
2950                reconnect_supported: false,
2951            },
2952        ))
2953    }
2954
2955    /// Returns a builder for a websocket client in **handler mode**.
2956    ///
2957    /// The handler is called for each incoming message on an internal task.
2958    /// Automatic reconnection is **enabled** with exponential backoff. On disconnection,
2959    /// the client automatically attempts to reconnect and replaces the internal reader
2960    /// (the handler continues working seamlessly).
2961    ///
2962    /// Use handler mode for simplified connection management, automatic reconnection, or
2963    /// callback-based message handling.
2964    ///
2965    /// See [`WebSocketConfig`] documentation for comparison with stream mode.
2966    ///
2967    /// Set `rate_limiter` to share message quota state across clients. Otherwise, the client
2968    /// creates one from `default_quota` and `keyed_quotas`. `connection_rate_limiter` gates the
2969    /// initial connection and reconnects using `connection_rate_keys`.
2970    ///
2971    /// Without `initial_connect_retry_policy` the builder makes exactly one connection attempt.
2972    /// With one, failures classified as retryable are retried up to its `max_attempts`; see
2973    /// [`InitialConnectRetryPolicy`] for which failures return before that bound is reached.
2974    ///
2975    /// `cancellation_token` aborts the initial connection only, and is observed during the
2976    /// connection rate-limit wait, the dial itself, and each backoff delay. It has no effect once
2977    /// this function returns a client: use [`WebSocketClient::disconnect`] to stop an established one, whose
2978    /// reconnect loop the token does not govern.
2979    ///
2980    /// Set `close_handler` to observe peer close frames and decide whether automatic reconnection
2981    /// continues. Close frames are not delivered to the message handler.
2982    ///
2983    /// The message handler is required:
2984    ///
2985    /// ```compile_fail
2986    /// use nautilus_network::websocket::{WebSocketClient, WebSocketConfig};
2987    ///
2988    /// let config: WebSocketConfig = unimplemented!();
2989    /// let _ = WebSocketClient::builder().config(config).connect();
2990    /// ```
2991    ///
2992    /// # Errors
2993    ///
2994    /// Returns an error if:
2995    /// - The configuration is invalid or the connection cannot be established.
2996    /// - A shared rate limiter is combined with quota configuration.
2997    /// - The connection rate limiter and its keys are not configured together.
2998    #[builder(finish_fn = connect)]
2999    pub async fn builder(
3000        config: WebSocketConfig,
3001        message_handler: MessageHandler,
3002        ping_handler: Option<PingHandler>,
3003        close_handler: Option<CloseHandler>,
3004        #[builder(default)] keyed_quotas: Vec<(String, Quota)>,
3005        default_quota: Option<Quota>,
3006        rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3007        state_sink: Option<SocketStateSink>,
3008        connection_rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3009        #[builder(default)] connection_rate_keys: Arc<[Ustr]>,
3010        initial_connect_retry_policy: Option<InitialConnectRetryPolicy>,
3011        cancellation_token: Option<CancellationToken>,
3012    ) -> Result<Self, TransportError> {
3013        let rate_limiter = Self::resolve_rate_limiter(default_quota, keyed_quotas, rate_limiter)?;
3014        let connection_rate_limit =
3015            Self::resolve_connection_rate_limit(connection_rate_limiter, connection_rate_keys)?;
3016        Self::connect_with_handler_scoped(
3017            config,
3018            IncomingHandler::Message(message_handler),
3019            ping_handler.map(IncomingPingHandler::Ping),
3020            close_handler,
3021            rate_limiter,
3022            state_sink,
3023            connection_rate_limit,
3024            InitialConnectOptions {
3025                retry_policy: initial_connect_retry_policy,
3026                cancellation_token,
3027            },
3028        )
3029        .await
3030    }
3031
3032    /// Returns a builder for a handler-mode client whose messages carry connection ownership.
3033    ///
3034    /// The initial connection has epoch `0`. Each replacement connection increments the epoch,
3035    /// and both its incoming messages and `RECONNECTED` notification carry that new value. Use
3036    /// [`Self::send_text_on_connection`] to bind an outgoing message to one of those epochs.
3037    /// Rate-limit, state, initial-connect retry, and cancellation options match [`Self::builder`].
3038    /// Set either `ping_handler` or `epoch_ping_handler` when custom ping handling is required.
3039    /// `close_handler` has the same behavior as [`Self::builder`].
3040    ///
3041    /// The epoch handler is required:
3042    ///
3043    /// ```compile_fail
3044    /// use nautilus_network::websocket::{WebSocketClient, WebSocketConfig};
3045    ///
3046    /// let config: WebSocketConfig = unimplemented!();
3047    /// let _ = WebSocketClient::epoch_builder().config(config).connect();
3048    /// ```
3049    ///
3050    /// # Errors
3051    ///
3052    /// Returns an error if:
3053    /// - The configuration is invalid or the connection cannot be established.
3054    /// - A shared rate limiter is combined with quota configuration.
3055    /// - The connection rate limiter and its keys are not configured together.
3056    /// - Both `ping_handler` and `epoch_ping_handler` are configured.
3057    #[builder(
3058        builder_type = WebSocketClientEpochBuilder,
3059        finish_fn = connect
3060    )]
3061    pub async fn epoch_builder(
3062        config: WebSocketConfig,
3063        epoch_handler: EpochMessageHandler,
3064        ping_handler: Option<PingHandler>,
3065        epoch_ping_handler: Option<EpochPingHandler>,
3066        close_handler: Option<CloseHandler>,
3067        #[builder(default)] keyed_quotas: Vec<(String, Quota)>,
3068        default_quota: Option<Quota>,
3069        rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3070        state_sink: Option<SocketStateSink>,
3071        connection_rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3072        #[builder(default)] connection_rate_keys: Arc<[Ustr]>,
3073        initial_connect_retry_policy: Option<InitialConnectRetryPolicy>,
3074        cancellation_token: Option<CancellationToken>,
3075    ) -> Result<Self, TransportError> {
3076        let ping_handler = match (ping_handler, epoch_ping_handler) {
3077            (Some(_), Some(_)) => {
3078                return Err(TransportError::Io(std::io::Error::new(
3079                    std::io::ErrorKind::InvalidInput,
3080                    "Cannot configure both ping_handler and epoch_ping_handler",
3081                )));
3082            }
3083            (Some(handler), None) => Some(IncomingPingHandler::Ping(handler)),
3084            (None, Some(handler)) => Some(IncomingPingHandler::Epoch(handler)),
3085            (None, None) => None,
3086        };
3087
3088        let rate_limiter = Self::resolve_rate_limiter(default_quota, keyed_quotas, rate_limiter)?;
3089        let connection_rate_limit =
3090            Self::resolve_connection_rate_limit(connection_rate_limiter, connection_rate_keys)?;
3091        Self::connect_with_handler_scoped(
3092            config,
3093            IncomingHandler::Epoch(epoch_handler),
3094            ping_handler,
3095            close_handler,
3096            rate_limiter,
3097            state_sink,
3098            connection_rate_limit,
3099            InitialConnectOptions {
3100                retry_policy: initial_connect_retry_policy,
3101                cancellation_token,
3102            },
3103        )
3104        .await
3105    }
3106
3107    fn resolve_connection_rate_limit(
3108        rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3109        keys: Arc<[Ustr]>,
3110    ) -> Result<Option<ConnectionRateLimit>, TransportError> {
3111        if rate_limiter.is_none() && !keys.is_empty() {
3112            return Err(TransportError::Io(std::io::Error::new(
3113                std::io::ErrorKind::InvalidInput,
3114                "Connection rate keys require a connection rate limiter",
3115            )));
3116        }
3117
3118        if rate_limiter.is_some() && keys.is_empty() {
3119            return Err(TransportError::Io(std::io::Error::new(
3120                std::io::ErrorKind::InvalidInput,
3121                "Connection rate limiter requires at least one connection rate key",
3122            )));
3123        }
3124
3125        Ok(rate_limiter.map(|limiter| ConnectionRateLimit { limiter, keys }))
3126    }
3127
3128    fn resolve_rate_limiter(
3129        default_quota: Option<Quota>,
3130        keyed_quotas: Vec<(String, Quota)>,
3131        rate_limiter: Option<Arc<RateLimiter<Ustr, MonotonicClock>>>,
3132    ) -> Result<Arc<RateLimiter<Ustr, MonotonicClock>>, TransportError> {
3133        if let Some(rate_limiter) = rate_limiter {
3134            if default_quota.is_some() || !keyed_quotas.is_empty() {
3135                return Err(TransportError::Io(std::io::Error::new(
3136                    std::io::ErrorKind::InvalidInput,
3137                    "Cannot combine a shared rate limiter with quota configuration",
3138                )));
3139            }
3140
3141            return Ok(rate_limiter);
3142        }
3143
3144        let keyed_quotas = keyed_quotas
3145            .into_iter()
3146            .map(|(key, quota)| (Ustr::from(&key), quota))
3147            .collect();
3148        Ok(Arc::new(RateLimiter::new_with_quota(
3149            default_quota,
3150            keyed_quotas,
3151        )))
3152    }
3153
3154    #[expect(
3155        clippy::too_many_arguments,
3156        reason = "handler delivery and connection controls are independent options"
3157    )]
3158    async fn connect_with_handler_scoped(
3159        config: WebSocketConfig,
3160        handler: IncomingHandler,
3161        ping_handler: Option<IncomingPingHandler>,
3162        close_handler: Option<CloseHandler>,
3163        rate_limiter: Arc<RateLimiter<Ustr, MonotonicClock>>,
3164        state_sink: Option<SocketStateSink>,
3165        connection_rate_limit: Option<ConnectionRateLimit>,
3166        initial_connect_options: InitialConnectOptions,
3167    ) -> Result<Self, TransportError> {
3168        log::debug!("Connecting");
3169        let inner = WebSocketClientInner::connect_url_with_handler(
3170            config,
3171            Some(handler),
3172            ping_handler,
3173            close_handler,
3174            state_sink,
3175            connection_rate_limit,
3176            initial_connect_options,
3177        )
3178        .await?;
3179        let connection_mode = Arc::clone(&inner.connection_mode);
3180        let connection_epoch = Arc::clone(&inner.connection_epoch);
3181        let state_notify = Arc::clone(&inner.state_notify);
3182        let controller_notify = Arc::clone(&inner.controller_notify);
3183        let reconnect_published = Arc::clone(&inner.reconnect_published);
3184        let writer_tx = inner.writer_tx.clone();
3185        let connect_timeout = inner.connect_timeout;
3186        let auth_tracker = Arc::clone(&inner.auth_tracker);
3187        let reconnect_buffer_waits_for_auth = Arc::clone(&inner.reconnect_buffer_waits_for_auth);
3188        let reconnect_headers = inner.reconnect_headers.clone();
3189        let state_sink = inner.state_sink.clone();
3190        let controller_lifecycle = Arc::new(ControllerLifecycle::new());
3191
3192        let controller_task = Self::spawn_controller_task(
3193            inner,
3194            Arc::clone(&connection_mode),
3195            Arc::clone(&state_notify),
3196            Arc::clone(&auth_tracker),
3197            Arc::clone(&controller_lifecycle),
3198            Arc::clone(&controller_notify),
3199            Arc::clone(&reconnect_published),
3200        );
3201        let abort_handle = controller_task.abort_handle();
3202        controller_lifecycle.set_abort(move || abort_handle.abort());
3203
3204        Ok(Self {
3205            controller_task,
3206            connection_mode,
3207            connection_epoch,
3208            state_notify,
3209            connect_timeout,
3210            rate_limiter,
3211            writer_tx,
3212            auth_tracker,
3213            reconnect_buffer_waits_for_auth,
3214            reconnect_headers,
3215            state_sink,
3216            controller_lifecycle,
3217            controller_notify,
3218            reconnect_published,
3219            reconnect_supported: true,
3220        })
3221    }
3222
3223    /// Returns shared headers used by future automatic reconnect attempts.
3224    #[must_use]
3225    pub fn reconnect_headers(&self) -> ReconnectHeaders {
3226        self.reconnect_headers.clone()
3227    }
3228
3229    /// Returns a cloneable handle to this client's reconnect controller.
3230    #[must_use]
3231    pub fn reconnect_handle(&self) -> WebSocketReconnectHandle {
3232        WebSocketReconnectHandle {
3233            connection_mode: Arc::clone(&self.connection_mode),
3234            auth_tracker: Arc::clone(&self.auth_tracker),
3235            state_sink: self.state_sink.clone(),
3236            controller_lifecycle: Arc::clone(&self.controller_lifecycle),
3237            controller_notify: Arc::clone(&self.controller_notify),
3238            reconnect_published: Arc::clone(&self.reconnect_published),
3239            supported: self.reconnect_supported,
3240        }
3241    }
3242
3243    /// Requests that the controller replace the active transport.
3244    ///
3245    /// Returns `true` only when this call transitions a handler-mode client from active to
3246    /// reconnecting. Stream-mode, duplicate, disconnecting, and closed requests return `false`.
3247    #[must_use]
3248    pub fn request_reconnect(&self) -> bool {
3249        self.reconnect_handle().request_reconnect() == ReconnectRequestOutcome::Accepted
3250    }
3251
3252    /// Returns the current connection mode.
3253    #[must_use]
3254    pub fn connection_mode(&self) -> ConnectionMode {
3255        ConnectionMode::from_atomic(&self.connection_mode)
3256    }
3257
3258    /// Returns the ownership epoch of the current connection.
3259    ///
3260    /// A connection keeps one epoch for its lifetime. The value increments when the writer swaps
3261    /// to a replacement connection.
3262    #[must_use]
3263    pub fn connection_epoch(&self) -> u64 {
3264        self.connection_epoch.load(Ordering::Acquire)
3265    }
3266
3267    /// Returns a clone of the connection mode atomic for external state tracking.
3268    ///
3269    /// This allows adapter clients to track connection state across reconnections
3270    /// without message-passing delays.
3271    #[must_use]
3272    pub fn connection_mode_atomic(&self) -> Arc<AtomicU8> {
3273        Arc::clone(&self.connection_mode)
3274    }
3275
3276    /// Returns shared read access to the current connection epoch.
3277    ///
3278    /// Callers must treat the returned atomic as read-only. The WebSocket writer task exclusively
3279    /// advances it when installing a replacement connection.
3280    #[must_use]
3281    pub fn connection_epoch_atomic(&self) -> Arc<AtomicU64> {
3282        Arc::clone(&self.connection_epoch)
3283    }
3284
3285    /// Returns whether the client connection is active.
3286    ///
3287    /// Returns `true` if the client is connected and has not been signaled to disconnect.
3288    /// The client will automatically retry connection based on its configuration.
3289    #[inline]
3290    #[must_use]
3291    pub fn is_active(&self) -> bool {
3292        self.connection_mode().is_active()
3293    }
3294
3295    /// Returns whether the controller task has stopped.
3296    #[must_use]
3297    pub fn is_disconnected(&self) -> bool {
3298        self.controller_task.is_finished()
3299    }
3300
3301    /// Returns whether the client is reconnecting.
3302    ///
3303    /// Returns `true` if the client lost connection and is attempting to reestablish it.
3304    /// The client will automatically retry connection based on its configuration.
3305    #[inline]
3306    #[must_use]
3307    pub fn is_reconnecting(&self) -> bool {
3308        self.connection_mode().is_reconnect()
3309    }
3310
3311    /// Registers an [`AuthTracker`] with the client.
3312    ///
3313    /// When the controller detects a dead connection and transitions to
3314    /// `Reconnect`, it calls `invalidate()` on the tracker so that any
3315    /// pending authenticated sends see the state change immediately. Terminal
3316    /// transitions fail the tracker so pending auth waits can terminate.
3317    /// Set `reconnect_buffer_waits_for_auth` for clients that must not replay
3318    /// buffered messages until the next session authenticates.
3319    ///
3320    /// Call this once after construction, before any authenticated sends.
3321    pub fn set_auth_tracker(&self, tracker: AuthTracker, reconnect_buffer_waits_for_auth: bool) {
3322        let _ = self.auth_tracker.set(tracker);
3323        self.reconnect_buffer_waits_for_auth
3324            .store(reconnect_buffer_waits_for_auth, Ordering::Release);
3325    }
3326
3327    /// Returns whether the client is disconnecting.
3328    ///
3329    /// Returns `true` if the client is in disconnect mode.
3330    #[inline]
3331    #[must_use]
3332    pub fn is_disconnecting(&self) -> bool {
3333        self.connection_mode().is_disconnect()
3334    }
3335
3336    /// Returns whether the client is closed.
3337    ///
3338    /// Returns `true` if the client has been explicitly disconnected or reached
3339    /// maximum reconnection attempts. In this state, the client cannot be reused
3340    /// and a new client must be created for further connections.
3341    #[inline]
3342    #[must_use]
3343    pub fn is_closed(&self) -> bool {
3344        self.connection_mode().is_closed()
3345    }
3346
3347    /// Checks whether the connection is in a terminal state (disconnecting or closed).
3348    ///
3349    /// Single atomic load to fail fast before rate limiting or waiting.
3350    #[inline]
3351    fn check_not_terminal(&self) -> Result<(), SendError> {
3352        match self.connection_mode() {
3353            ConnectionMode::Disconnect | ConnectionMode::Closed => Err(SendError::Closed),
3354            _ => Ok(()),
3355        }
3356    }
3357
3358    /// Waits for rate limiter quota, aborting early if connection enters a terminal state.
3359    async fn await_rate_limit_or_closed(&self, keys: Option<&[Ustr]>) -> Result<(), SendError> {
3360        const CHECK_INTERVAL_MS: u64 = 100;
3361
3362        tokio::select! {
3363            biased;
3364            () = self.rate_limiter.await_keys_ready(keys) => Ok(()),
3365            () = async {
3366                loop {
3367                    // Enable before the state check: an unpolled Notified is unregistered and misses notifies
3368                    let mut notified = pin!(self.state_notify.notified());
3369                    notified.as_mut().enable();
3370
3371                    if matches!(self.connection_mode(), ConnectionMode::Disconnect | ConnectionMode::Closed) {
3372                        break;
3373                    }
3374                    tokio::select! {
3375                        biased;
3376                        () = notified => {}
3377                        () = dst::time::sleep(Duration::from_millis(CHECK_INTERVAL_MS)) => {}
3378                    }
3379                }
3380            } => Err(SendError::Closed),
3381        }
3382    }
3383
3384    /// Waits for the client to become active before sending.
3385    ///
3386    /// Uses `state_notify` for event-driven wakeup so sends resume immediately
3387    /// after reconnection completes. A fallback interval guards against missed
3388    /// notifications.
3389    async fn wait_for_active(&self) -> Result<(), SendError> {
3390        const FALLBACK_INTERVAL_MS: u64 = 100;
3391
3392        let mode = self.connection_mode();
3393        if mode.is_active() {
3394            return Ok(());
3395        }
3396
3397        if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
3398            return Err(SendError::Closed);
3399        }
3400
3401        log::debug!("Waiting for client to become ACTIVE before sending...");
3402
3403        let fallback_interval = Duration::from_millis(FALLBACK_INTERVAL_MS);
3404
3405        dst::time::timeout(self.connect_timeout, async {
3406            loop {
3407                // Enable before the state check: an unpolled Notified is unregistered and misses notifies
3408                let mut notified = pin!(self.state_notify.notified());
3409                notified.as_mut().enable();
3410
3411                let mode = self.connection_mode();
3412                if mode.is_active() {
3413                    return Ok(());
3414                }
3415
3416                if matches!(mode, ConnectionMode::Disconnect | ConnectionMode::Closed) {
3417                    return Err(());
3418                }
3419
3420                tokio::select! {
3421                    biased;
3422                    () = notified => {}
3423                    () = dst::time::sleep(fallback_interval) => {}
3424                }
3425            }
3426        })
3427        .await
3428        .map_err(|_| SendError::Timeout)?
3429        .map_err(|()| SendError::Closed)
3430    }
3431
3432    /// Signals that the caller's reader has observed EOF or a fatal error.
3433    ///
3434    /// In stream mode the controller has no visibility into the caller-owned reader.
3435    /// Call this method when `reader.next().await` returns `None` or an unrecoverable
3436    /// error so the controller transitions to `Closed` and dependent tasks shut down.
3437    ///
3438    /// For peer-initiated close frames (`Message::Close`), use [`disconnect`](Self::disconnect)
3439    /// instead so the writer can send the close reply before shutting down.
3440    ///
3441    /// If an [`AuthTracker`] is registered, this fails pending auth waits.
3442    ///
3443    /// This is a no-op if the connection is already closed or disconnecting.
3444    pub fn notify_closed(&self) {
3445        let mode = self.connection_mode();
3446        if mode.is_disconnect() || mode.is_closed() {
3447            return;
3448        }
3449
3450        log::debug!("Stream reader signalled EOF, transitioning to CLOSED");
3451
3452        if ConnectionMode::close_websocket_on_loss(&self.connection_mode, self.state_sink.as_ref())
3453        {
3454            fail_registered_auth(self.auth_tracker.as_ref(), "WebSocket client closed");
3455            self.state_notify.notify_waiters();
3456        }
3457    }
3458
3459    /// Disconnects the client and waits for the controller task to stop.
3460    ///
3461    /// If an [`AuthTracker`] is registered, this fails pending auth waits.
3462    pub async fn disconnect(&self) {
3463        log::debug!("Disconnecting");
3464
3465        // A CLOSED client keeps its terminal state; its tracker is already failed
3466        if ConnectionMode::request_disconnect(&self.connection_mode)
3467            && let Some(tracker) = self.auth_tracker.get()
3468        {
3469            tracker.fail("WebSocket client disconnected");
3470        }
3471
3472        self.state_notify.notify_waiters();
3473
3474        if dst::time::timeout(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS), async {
3475            while !self.is_disconnected() {
3476                dst::time::sleep(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
3477            }
3478
3479            if !self.controller_task.is_finished() {
3480                self.controller_task.abort();
3481                log_task_aborted("controller");
3482            }
3483        })
3484        .await
3485            == Ok(())
3486        {
3487            log::debug!("Controller task finished");
3488        } else {
3489            log::warn!("Timeout waiting for controller task to finish");
3490
3491            if !self.controller_task.is_finished() {
3492                self.controller_task.abort();
3493                log_task_aborted("controller");
3494            }
3495
3496            self.connection_mode
3497                .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3498        }
3499    }
3500
3501    /// Sends the given text `data` to the server.
3502    ///
3503    /// Returns `Ok(())` when the message is enqueued to the writer channel. This does not
3504    /// guarantee delivery: if a disconnect occurs concurrently, the writer task may drop the
3505    /// message. During reconnection, messages are buffered and replayed on the new connection.
3506    ///
3507    /// # Errors
3508    ///
3509    /// Returns an error if the writer capacity is exhausted, or if the client cannot send.
3510    #[allow(unused_variables)]
3511    pub async fn send_text(&self, data: String, keys: Option<&[Ustr]>) -> Result<(), SendError> {
3512        self.check_not_terminal()?;
3513
3514        self.await_rate_limit_or_closed(keys).await?;
3515        self.wait_for_active().await?;
3516
3517        log::trace!("Sending text frame ({} bytes)", data.len());
3518
3519        let msg = Message::Text(data.into());
3520        self.writer_tx.send(WriterCommand::Send(msg))
3521    }
3522
3523    /// Sends text once if the active connection matches `connection_epoch`.
3524    ///
3525    /// The writer rejects the send if reconnection changes ownership before the write. The
3526    /// message is never replayed on another connection.
3527    ///
3528    /// # Errors
3529    ///
3530    /// Returns an error if:
3531    /// - The shared control allowance is full.
3532    /// - The client closes.
3533    /// - The active wait times out before the write starts.
3534    /// - The expected connection no longer owns the writer.
3535    /// - The command or transport write fails.
3536    /// - The write starts but does not complete within the write
3537    ///   deadline. Delivery is undetermined in that case: the message is not replayed, and it
3538    ///   must not be resent blindly.
3539    pub async fn send_text_on_connection(
3540        &self,
3541        data: String,
3542        keys: Option<&[Ustr]>,
3543        connection_epoch: u64,
3544    ) -> Result<(), SendError> {
3545        self.check_not_terminal()?;
3546        self.await_rate_limit_or_closed(keys).await?;
3547        self.wait_for_active().await?;
3548
3549        log::trace!(
3550            "Sending text frame once: epoch={connection_epoch} ({} bytes)",
3551            data.len()
3552        );
3553
3554        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
3555        self.writer_tx
3556            .send_control(WriterCommand::SendOnConnection {
3557                message: Message::Text(data.into()),
3558                connection_epoch,
3559                response_tx,
3560            })?;
3561
3562        response_rx
3563            .await
3564            .map_err(|e| SendError::BrokenPipe(e.to_string()))?
3565    }
3566
3567    /// Sends a pong frame back to the server when the connection is active.
3568    ///
3569    /// The pong is skipped silently if the connection is not active when called:
3570    /// a pong belongs to the connection whose ping caused it, so this method does
3571    /// not wait for reconnection before enqueueing it.
3572    ///
3573    /// # Errors
3574    ///
3575    /// Returns an error if:
3576    /// - The payload exceeds 125 bytes, the RFC 6455 control-frame limit.
3577    /// - The writer channel is broken or its control allowance is full.
3578    #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
3579    #[expect(
3580        clippy::unused_async,
3581        clippy::unused_async_trait_impl,
3582        reason = "skipping instead of waiting removes the only await; the signature is public API shared with the other send methods"
3583    )]
3584    pub async fn send_pong(&self, data: Vec<u8>) -> Result<(), SendError> {
3585        validate_pong_payload(&data)?;
3586
3587        if !self.connection_mode().is_active() {
3588            log::debug!("Skipping pong: connection not active");
3589            return Ok(());
3590        }
3591
3592        log::trace!("Sending pong frame ({} bytes)", data.len());
3593
3594        self.writer_tx
3595            .send_control(WriterCommand::Send(Message::Pong(data.into())))
3596    }
3597
3598    /// Sends a pong if the connection that received its ping is still active.
3599    ///
3600    /// The pong is skipped silently if the connection is inactive or its epoch has changed.
3601    ///
3602    /// # Errors
3603    ///
3604    /// Returns an error if:
3605    /// - The payload exceeds 125 bytes, the RFC 6455 control-frame limit.
3606    /// - The writer channel is broken or its control allowance is full.
3607    #[allow(unknown_lints, reason = "Clippy lint is unavailable on Rust 1.97")]
3608    #[expect(
3609        clippy::unused_async,
3610        clippy::unused_async_trait_impl,
3611        reason = "the public send API is async even though this method only enqueues"
3612    )]
3613    pub async fn send_pong_on_connection(
3614        &self,
3615        data: Vec<u8>,
3616        connection_epoch: u64,
3617    ) -> Result<(), SendError> {
3618        validate_pong_payload(&data)?;
3619
3620        if !self.connection_mode().is_active() {
3621            log::debug!("Skipping pong: connection not active");
3622            return Ok(());
3623        }
3624
3625        log::trace!(
3626            "Sending pong frame once: epoch={connection_epoch} ({} bytes)",
3627            data.len()
3628        );
3629
3630        self.writer_tx
3631            .send_control(WriterCommand::SendPongOnConnection {
3632                data,
3633                connection_epoch,
3634            })
3635    }
3636
3637    /// Sends the given bytes `data` to the server.
3638    ///
3639    /// Returns `Ok(())` when the message is enqueued to the writer channel. This does not
3640    /// guarantee delivery: if a disconnect occurs concurrently, the writer task may drop the
3641    /// message. During reconnection, messages are buffered and replayed on the new connection.
3642    ///
3643    /// # Errors
3644    ///
3645    /// Returns an error if the writer capacity is exhausted, or if the client cannot send.
3646    #[allow(unused_variables)]
3647    pub async fn send_bytes(&self, data: Vec<u8>, keys: Option<&[Ustr]>) -> Result<(), SendError> {
3648        self.check_not_terminal()?;
3649
3650        self.await_rate_limit_or_closed(keys).await?;
3651        self.wait_for_active().await?;
3652
3653        log::trace!("Sending binary frame ({} bytes)", data.len());
3654
3655        let msg = Message::Binary(data.into());
3656        self.writer_tx.send(WriterCommand::Send(msg))
3657    }
3658
3659    /// Sends a close message to the server.
3660    ///
3661    /// # Errors
3662    ///
3663    /// Returns an error if the shared control allowance is full, or if the client cannot send.
3664    pub async fn send_close_message(&self) -> Result<(), SendError> {
3665        self.wait_for_active().await?;
3666
3667        let msg = Message::Close(None);
3668        self.writer_tx.send_control(WriterCommand::Send(msg))
3669    }
3670
3671    fn spawn_controller_task(
3672        mut inner: WebSocketClientInner,
3673        connection_mode: Arc<AtomicU8>,
3674        state_notify: Arc<tokio::sync::Notify>,
3675        auth_tracker: Arc<OnceLock<AuthTracker>>,
3676        controller_lifecycle: Arc<ControllerLifecycle>,
3677        controller_notify: Arc<tokio::sync::Notify>,
3678        reconnect_published: Arc<AtomicBool>,
3679    ) -> dst::task::JoinHandle<()> {
3680        dst::task::spawn(async move {
3681            let _activity = controller_lifecycle.activity();
3682            log_task_started("controller");
3683
3684            let fallback_interval = Duration::from_millis(CONTROLLER_FALLBACK_INTERVAL_MS);
3685            let mut reconnected_at = None;
3686
3687            loop {
3688                tokio::select! {
3689                    biased;
3690                    () = controller_notify.notified() => {}
3691                    () = state_notify.notified() => {}
3692                    () = dst::time::sleep(fallback_interval) => {}
3693                }
3694
3695                let mut mode = ConnectionMode::from_atomic(&connection_mode);
3696
3697                if mode.is_disconnect() {
3698                    log::debug!("Disconnecting");
3699
3700                    let timeout = Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS);
3701                    if dst::time::timeout(timeout, async {
3702                        // Delay awaiting graceful shutdown
3703                        dst::time::sleep(Duration::from_millis(GRACEFUL_SHUTDOWN_DELAY_MS)).await;
3704
3705                        if let Some(read_fence) = inner.read_fence.take() {
3706                            read_fence.invalidate();
3707                        }
3708
3709                        if let Some(task) = &inner.read_task
3710                            && !task.is_finished()
3711                        {
3712                            task.abort();
3713                            log_task_aborted("read");
3714                        }
3715
3716                        if let Some(task) = &inner.heartbeat_task
3717                            && !task.is_finished()
3718                        {
3719                            task.abort();
3720                            log_task_aborted("heartbeat");
3721                        }
3722                    })
3723                    .await
3724                    .is_err()
3725                    {
3726                        log::warn!("Shutdown timed out after {}s", timeout.as_secs());
3727                    }
3728
3729                    log::debug!("Closed");
3730                    break; // Controller finished
3731                }
3732
3733                if mode.is_closed() {
3734                    log::debug!("Connection closed");
3735                    fail_registered_auth(auth_tracker.as_ref(), "WebSocket client closed");
3736                    state_notify.notify_waiters();
3737                    break;
3738                }
3739
3740                if mode.is_active() && !inner.is_alive() {
3741                    let target = if inner.handler.is_none() {
3742                        ConnectionMode::Closed
3743                    } else {
3744                        ConnectionMode::Reconnect
3745                    };
3746
3747                    let transitioned = if target.is_closed() {
3748                        ConnectionMode::close_websocket_on_loss(
3749                            &connection_mode,
3750                            inner.state_sink.as_ref(),
3751                        )
3752                    } else {
3753                        request_websocket_reconnect(
3754                            &connection_mode,
3755                            &reconnect_published,
3756                            inner.state_sink.as_ref(),
3757                            &auth_tracker,
3758                            &controller_notify,
3759                            || {},
3760                        ) == ReconnectRequestOutcome::Accepted
3761                    };
3762
3763                    if transitioned {
3764                        if target.is_closed() {
3765                            fail_registered_auth(auth_tracker.as_ref(), "WebSocket client closed");
3766                        }
3767
3768                        log::info!("Detected dead connection, transitioning to {target:?}");
3769                    }
3770
3771                    mode = ConnectionMode::from_atomic(&connection_mode);
3772                }
3773
3774                if mode.is_reconnect() {
3775                    if let Some(tracker) = auth_tracker.get() {
3776                        tracker.invalidate();
3777                    }
3778
3779                    let reconnect_uptime = reconnected_at
3780                        .take()
3781                        .map(|started: dst::time::Instant| started.elapsed());
3782                    let previous_reconnect_stable = reconnect_uptime
3783                        .is_some_and(|uptime| uptime >= RECONNECT_STABILITY_THRESHOLD);
3784
3785                    if previous_reconnect_stable {
3786                        inner.backoff.reset();
3787                        inner.reconnection_attempt_count = 0;
3788                        log::debug!(
3789                            "WebSocket remained active for at least {}s, resetting reconnect cycle",
3790                            RECONNECT_STABILITY_THRESHOLD.as_secs()
3791                        );
3792                    }
3793
3794                    if let Some(max_attempts) = inner.reconnect_max_attempts
3795                        && inner.reconnection_attempt_count >= max_attempts
3796                    {
3797                        log::error!(
3798                            "Max reconnection attempts ({max_attempts}) exceeded, transitioning to CLOSED"
3799                        );
3800                        connection_mode.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3801                        fail_registered_auth(
3802                            auth_tracker.as_ref(),
3803                            "WebSocket reconnect attempts exhausted",
3804                        );
3805                        state_notify.notify_waiters();
3806                        break;
3807                    }
3808
3809                    let backoff_delay = if reconnect_uptime.is_some() && !previous_reconnect_stable
3810                    {
3811                        inner.backoff.next_duration()
3812                    } else {
3813                        Duration::ZERO
3814                    };
3815
3816                    let duration = inner.reconnect_throttle.gated_delay(backoff_delay);
3817                    if !duration.is_zero() {
3818                        log::warn!("Backing off for {}s...", duration.as_secs_f64());
3819
3820                        if !wait_reconnect_delay(
3821                            duration,
3822                            connection_mode.as_ref(),
3823                            state_notify.as_ref(),
3824                        )
3825                        .await
3826                        {
3827                            log::debug!("Backoff interrupted by terminal state");
3828                            continue;
3829                        }
3830                    }
3831
3832                    inner.reconnection_attempt_count += 1;
3833                    inner.reconnect_throttle.record_attempt();
3834                    log::debug!(
3835                        "Reconnection attempt {} of {}",
3836                        inner.reconnection_attempt_count,
3837                        inner
3838                            .reconnect_max_attempts
3839                            .map_or_else(|| "unlimited".to_string(), |m| m.to_string())
3840                    );
3841
3842                    // Race reconnect against disconnect notification
3843                    let reconnect_result = tokio::select! {
3844                        biased;
3845                        result = inner.reconnect_with_outcome() => Some(result),
3846                        () = async {
3847                            loop {
3848                                // Enable before the check so a disconnect notify between iterations is not missed
3849                                let mut notified = pin!(state_notify.notified());
3850                                notified.as_mut().enable();
3851
3852                                if ConnectionMode::from_atomic(&connection_mode).is_disconnect() {
3853                                    break;
3854                                }
3855                                notified.await;
3856                            }
3857                        } => None,
3858                    };
3859
3860                    match reconnect_result {
3861                        None => {
3862                            log::debug!("Reconnect interrupted by disconnect");
3863                        }
3864                        Some(Ok(ReconnectOutcome::Reconnected)) => {
3865                            reconnected_at = Some(dst::time::Instant::now());
3866
3867                            state_notify.notify_waiters();
3868
3869                            // The outcome records a completed reconnection; emit recovery
3870                            // callbacks only while the replacement is still `Active`, not
3871                            // after a teardown or another drop.
3872                            if ConnectionMode::from_atomic(&connection_mode).is_active() {
3873                                if let Some(ref handler) = inner.handler {
3874                                    let connection_epoch =
3875                                        inner.connection_epoch.load(Ordering::Acquire);
3876                                    let reconnected_msg =
3877                                        Message::Text(RECONNECTED.to_string().into());
3878                                    handler.handle(connection_epoch, reconnected_msg);
3879                                    match handler {
3880                                        IncomingHandler::Message(_) => {
3881                                            log::debug!("Sent reconnected message to handler");
3882                                        }
3883                                        IncomingHandler::Epoch(_) => {
3884                                            log::debug!(
3885                                                "Sent reconnected message to epoch handler: \
3886                                                 epoch={connection_epoch}",
3887                                            );
3888                                        }
3889                                    }
3890                                }
3891
3892                                log::debug!("Reconnected successfully");
3893                            } else {
3894                                log::debug!("Skipping reconnect handlers due to disconnect state");
3895                            }
3896                        }
3897                        Some(Ok(ReconnectOutcome::Aborted)) => {
3898                            log::debug!("Reconnect aborted");
3899                        }
3900                        Some(Err(e)) => {
3901                            let duration = inner.backoff.next_duration();
3902                            log::warn!(
3903                                "Reconnect attempt {} failed: {e}",
3904                                inner.reconnection_attempt_count
3905                            );
3906
3907                            if !duration.is_zero() {
3908                                log::warn!("Backing off for {}s...", duration.as_secs_f64());
3909
3910                                if !wait_reconnect_delay(
3911                                    duration,
3912                                    connection_mode.as_ref(),
3913                                    state_notify.as_ref(),
3914                                )
3915                                .await
3916                                {
3917                                    log::debug!("Backoff interrupted by terminal state");
3918                                }
3919                            }
3920                        }
3921                    }
3922                }
3923            }
3924
3925            inner
3926                .connection_mode
3927                .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3928
3929            log_task_stopped("controller");
3930        })
3931    }
3932}
3933
3934fn fail_registered_auth(auth_tracker: &OnceLock<AuthTracker>, reason: &str) {
3935    if let Some(tracker) = auth_tracker.get() {
3936        tracker.fail(reason);
3937    }
3938}
3939
3940fn validate_pong_payload(data: &[u8]) -> Result<(), SendError> {
3941    if data.len() > MAX_CONTROL_FRAME_PAYLOAD_BYTES {
3942        return Err(SendError::InvalidInput(format!(
3943            "pong payload exceeds {MAX_CONTROL_FRAME_PAYLOAD_BYTES} bytes"
3944        )));
3945    }
3946
3947    Ok(())
3948}
3949
3950impl Drop for WebSocketClient {
3951    fn drop(&mut self) {
3952        self.connection_mode
3953            .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
3954        fail_registered_auth(self.auth_tracker.as_ref(), "WebSocket client closed");
3955        self.state_notify.notify_waiters();
3956        self.controller_notify.notify_waiters();
3957        self.controller_lifecycle.close_and_abort();
3958    }
3959}
3960
3961#[cfg(test)]
3962#[cfg(not(feature = "turmoil"))]
3963#[cfg(not(all(feature = "simulation", madsim)))] // These tests drive real localhost sockets
3964#[cfg(target_os = "linux")] // Only run network tests on Linux (CI stability)
3965mod tests {
3966    use std::{
3967        collections::HashMap,
3968        num::NonZeroU32,
3969        sync::{
3970            Arc,
3971            atomic::{AtomicU8, AtomicUsize, Ordering},
3972        },
3973        time::Duration,
3974    };
3975
3976    use axum::{Router, routing::post};
3977    use futures_util::{SinkExt, StreamExt};
3978    use log::Level;
3979    use nautilus_common::testing::wait_until_async;
3980    use nautilus_core::string::secret::REDACTED;
3981    use parking_lot::Mutex;
3982    use rstest::rstest;
3983    use tokio::{
3984        io::{AsyncReadExt, AsyncWriteExt},
3985        net::TcpListener,
3986        sync::{mpsc, oneshot},
3987        task::{self, JoinHandle},
3988    };
3989    use tokio_tungstenite::{
3990        accept_async, accept_hdr_async,
3991        tungstenite::{
3992            Message as WsMessage,
3993            handshake::server::{self, Callback},
3994            http::HeaderValue,
3995            protocol::frame::{CloseFrame as WsCloseFrame, coding::CloseCode},
3996        },
3997    };
3998
3999    use crate::{
4000        RECONNECTED, SocketState, SocketStateSink,
4001        error::SendError,
4002        http::{HttpClient, Method},
4003        logging::tests::capture_logs_for,
4004        mode::{ConnectionMode, ReadSessionFence},
4005        ratelimiter::quota::Quota,
4006        transport::{CloseFrame, Message, TransportError},
4007        websocket::{
4008            AuthTracker, CloseHandler, InitialConnectRetryPolicy, ReconnectHeaders,
4009            TransportBackend, WebSocketClient, WebSocketClientInner, WebSocketConfig,
4010            auth::AuthState,
4011            types::{WriterCommand, channel_epoch_message_handler, channel_message_handler},
4012        },
4013    };
4014
4015    const SECRET_MARKER: &str = "OUTBOUND_SECRET_MARKER";
4016    const PING_TRIGGER: &str = "send-test-ping";
4017    const HOSTILE_CLOSE_TRIGGER: &str = "close-hostile";
4018    const CODED_CLOSE_TRIGGER: &str = "close-coded";
4019    const CODED_CLOSE_CODE: u16 = 4003;
4020    const CODED_CLOSE_REASON: &str = "draining";
4021    const HOSTILE_CLOSE_REASON: &str = "injected\r\n\u{1b}[31mforged\u{1b}[0m";
4022    const NETWORK_LOG_TARGETS: &[&str] = &[
4023        "nautilus_network::http::client",
4024        "nautilus_network::websocket::client",
4025    ];
4026
4027    struct TestServer {
4028        task: JoinHandle<()>,
4029        port: u16,
4030    }
4031
4032    #[derive(Debug, Clone)]
4033    struct TestCallback {
4034        key: String,
4035        value: HeaderValue,
4036    }
4037
4038    impl Callback for TestCallback {
4039        #[expect(clippy::panic_in_result_fn)]
4040        fn on_request(
4041            self,
4042            request: &server::Request,
4043            response: server::Response,
4044        ) -> Result<server::Response, server::ErrorResponse> {
4045            let _ = response;
4046            let value = request.headers().get(&self.key);
4047            assert!(value.is_some());
4048
4049            if let Some(value) = request.headers().get(&self.key) {
4050                assert_eq!(value, self.value);
4051            }
4052
4053            Ok(response)
4054        }
4055    }
4056
4057    impl TestServer {
4058        async fn setup() -> Self {
4059            let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
4060            let port = TcpListener::local_addr(&server).unwrap().port();
4061
4062            let header_key = "test".to_string();
4063            let header_value = "test".to_string();
4064
4065            let test_call_back = TestCallback {
4066                key: header_key,
4067                value: HeaderValue::from_str(&header_value).unwrap(),
4068            };
4069
4070            let task = task::spawn(async move {
4071                // Keep accepting connections
4072                loop {
4073                    let (conn, _) = server.accept().await.unwrap();
4074                    let mut websocket = accept_hdr_async(conn, test_call_back.clone())
4075                        .await
4076                        .unwrap();
4077
4078                    task::spawn(async move {
4079                        while let Some(Ok(msg)) = websocket.next().await {
4080                            match msg {
4081                                WsMessage::Text(txt) if txt == "close-now" => {
4082                                    log::debug!("Forcibly closing from server side");
4083                                    // This sends a close frame, then stops reading
4084                                    let _ = websocket.close(None).await;
4085                                    break;
4086                                }
4087                                WsMessage::Text(txt) if txt == HOSTILE_CLOSE_TRIGGER => {
4088                                    // Close reasons are server-controlled: this one carries a
4089                                    // forged line break and ANSI color sequences for the log
4090                                    // sanitization test.
4091                                    let _ = websocket
4092                                        .send(WsMessage::Close(Some(WsCloseFrame {
4093                                            code: CloseCode::Error,
4094                                            reason: HOSTILE_CLOSE_REASON.into(),
4095                                        })))
4096                                        .await;
4097
4098                                    break;
4099                                }
4100                                WsMessage::Text(txt) if txt == CODED_CLOSE_TRIGGER => {
4101                                    let _ = websocket
4102                                        .send(WsMessage::Close(Some(WsCloseFrame {
4103                                            code: CloseCode::from(CODED_CLOSE_CODE),
4104                                            reason: CODED_CLOSE_REASON.into(),
4105                                        })))
4106                                        .await;
4107
4108                                    break;
4109                                }
4110                                WsMessage::Text(txt) if txt == PING_TRIGGER => {
4111                                    let ping = format!("{SECRET_MARKER}:ping");
4112                                    if websocket.send(WsMessage::Ping(ping.into())).await.is_err() {
4113                                        break;
4114                                    }
4115                                }
4116                                // Echo text/binary frames
4117                                WsMessage::Text(_) | WsMessage::Binary(_) => {
4118                                    if websocket.send(msg).await.is_err() {
4119                                        break;
4120                                    }
4121                                }
4122                                // If the client closes, we also break
4123                                WsMessage::Close(_frame) => {
4124                                    let _ = websocket.close(None).await;
4125                                    break;
4126                                }
4127                                // Ignore pings/pongs
4128                                _ => {}
4129                            }
4130                        }
4131                    });
4132                }
4133            });
4134
4135            Self { task, port }
4136        }
4137    }
4138
4139    impl Drop for TestServer {
4140        fn drop(&mut self) {
4141            self.task.abort();
4142        }
4143    }
4144
4145    async fn setup_test_client(port: u16) -> WebSocketClient {
4146        let config = WebSocketConfig {
4147            url: format!("ws://127.0.0.1:{port}"),
4148            headers: vec![("test".into(), "test".into())],
4149            heartbeat_interval_secs: None,
4150            heartbeat_payload: None,
4151            connect_timeout_ms: None,
4152            reconnect_delay_initial_ms: None,
4153            reconnect_backoff_factor: None,
4154            reconnect_delay_max_ms: None,
4155            reconnect_jitter_ms: None,
4156            reconnect_max_attempts: None,
4157            heartbeat_timeout_secs: None,
4158            idle_timeout_ms: None,
4159            writer_capacity: None,
4160            backend: TransportBackend::Tungstenite,
4161            proxy_url: None,
4162            max_message_size_bytes: None,
4163            max_frame_size_bytes: None,
4164        };
4165
4166        WebSocketClient::builder()
4167            .config(config)
4168            .message_handler(Arc::new(|_| {}))
4169            .connect()
4170            .await
4171            .expect("Failed to connect")
4172    }
4173
4174    #[rstest]
4175    #[case::handler(false)]
4176    #[case::stream(true)]
4177    #[tokio::test]
4178    async fn test_writer_capacity_from_config(#[case] stream_mode: bool) {
4179        let server = TestServer::setup().await;
4180        let config = WebSocketConfig::builder()
4181            .url(format!("ws://127.0.0.1:{}", server.port))
4182            .headers(vec![("test".to_string(), "test".to_string())])
4183            .backend(TransportBackend::Tungstenite)
4184            .writer_capacity(2)
4185            .build()
4186            .unwrap();
4187
4188        let (client, _reader) = if stream_mode {
4189            let (reader, client) = WebSocketClient::stream_builder()
4190                .config(config)
4191                .connect()
4192                .await
4193                .unwrap();
4194            (client, Some(reader))
4195        } else {
4196            let client = WebSocketClient::builder()
4197                .config(config)
4198                .message_handler(Arc::new(|_| {}))
4199                .connect()
4200                .await
4201                .unwrap();
4202            (client, None)
4203        };
4204
4205        // Ready sends do not yield on this current-thread runtime, keeping the queue occupied
4206        client.send_text("first".to_string(), None).await.unwrap();
4207        client.send_bytes(vec![2], None).await.unwrap();
4208        let overflow = client.send_text("overflow".to_string(), None).await;
4209
4210        assert_eq!(client.connection_mode(), ConnectionMode::Active);
4211        assert!(matches!(overflow, Err(SendError::BufferFull)));
4212        client.disconnect().await;
4213    }
4214
4215    #[rstest]
4216    #[tokio::test]
4217    async fn test_send_reports_writer_buffer_full() {
4218        let server = TestServer::setup().await;
4219        let mut client = setup_test_client(server.port).await;
4220        let (writer_tx, mut writer_rx) = crate::writer::channel(1);
4221        client.writer_tx = writer_tx;
4222
4223        client.send_text("first".to_string(), None).await.unwrap();
4224        let result = client.send_bytes(vec![2], None).await;
4225        assert!(matches!(result, Err(SendError::BufferFull)));
4226        let (command, permit) = writer_rx.recv().await.unwrap();
4227        assert!(
4228            matches!(command, WriterCommand::Send(message) if message == Message::text("first"))
4229        );
4230        assert!(matches!(
4231            client.send_text("overflow".to_string(), None).await,
4232            Err(SendError::BufferFull)
4233        ));
4234        drop(permit);
4235        client.send_bytes(vec![3], None).await.unwrap();
4236        let (command, _) = writer_rx.recv().await.unwrap();
4237        assert!(
4238            matches!(command, WriterCommand::Send(message) if message == Message::Binary(vec![3].into()))
4239        );
4240        client.disconnect().await;
4241    }
4242
4243    async fn setup_reconnecting_client(port: u16) -> WebSocketClient {
4244        let config = WebSocketConfig {
4245            url: format!("ws://127.0.0.1:{port}"),
4246            headers: vec![],
4247            heartbeat_interval_secs: None,
4248            heartbeat_payload: None,
4249            connect_timeout_ms: Some(5_000),
4250            reconnect_delay_initial_ms: Some(1),
4251            reconnect_backoff_factor: Some(1.0),
4252            reconnect_delay_max_ms: Some(1),
4253            reconnect_jitter_ms: Some(0),
4254            reconnect_max_attempts: None,
4255            heartbeat_timeout_secs: None,
4256            idle_timeout_ms: None,
4257            writer_capacity: None,
4258            backend: TransportBackend::Tungstenite,
4259            proxy_url: None,
4260            max_message_size_bytes: None,
4261            max_frame_size_bytes: None,
4262        };
4263
4264        WebSocketClient::builder()
4265            .config(config)
4266            .message_handler(Arc::new(|_| {}))
4267            .connect()
4268            .await
4269            .expect("client should connect")
4270    }
4271
4272    async fn wait_for_mode(client: &WebSocketClient, expected: ConnectionMode) {
4273        crate::dst::time::timeout(Duration::from_secs(5), async {
4274            loop {
4275                if ConnectionMode::from_atomic(&client.connection_mode) == expected {
4276                    break;
4277                }
4278
4279                crate::dst::time::sleep(Duration::from_millis(1)).await;
4280            }
4281        })
4282        .await
4283        .expect("client should reach expected connection mode");
4284    }
4285
4286    async fn setup_http_test_server() -> (JoinHandle<()>, u16) {
4287        let server = TcpListener::bind("127.0.0.1:0").await.unwrap();
4288        let port = server.local_addr().unwrap().port();
4289
4290        let app = Router::new().route(
4291            "/logging",
4292            post(|| async {
4293                (
4294                    [("x-secret-response", SECRET_MARKER)],
4295                    format!("{SECRET_MARKER}:response"),
4296                )
4297            }),
4298        );
4299
4300        let task = task::spawn(async move {
4301            axum::serve(server, app).await.unwrap();
4302        });
4303
4304        (task, port)
4305    }
4306
4307    #[rstest]
4308    #[tokio::test]
4309    async fn test_network_logs_omit_payload_bodies() {
4310        let server = TestServer::setup().await;
4311        let client = setup_test_client(server.port).await;
4312        let capture = capture_logs_for(NETWORK_LOG_TARGETS).await;
4313        let binary = format!("{SECRET_MARKER}:binary").into_bytes();
4314        let binary_marker = format!("{binary:?}");
4315
4316        client
4317            .send_text(format!("{SECRET_MARKER}:café"), None)
4318            .await
4319            .unwrap();
4320        client
4321            .send_text_on_connection(
4322                format!("{SECRET_MARKER}:owned-é"),
4323                None,
4324                client.connection_epoch(),
4325            )
4326            .await
4327            .unwrap();
4328        client.send_bytes(binary, None).await.unwrap();
4329        client
4330            .send_text(PING_TRIGGER.to_string(), None)
4331            .await
4332            .unwrap();
4333
4334        tokio::time::timeout(Duration::from_secs(2), async {
4335            loop {
4336                if capture.messages().iter().any(|(level, message)| {
4337                    matches!(level, Level::Trace | Level::Warn)
4338                        && message == "Received ping frame (27 bytes)"
4339                }) {
4340                    break;
4341                }
4342
4343                tokio::task::yield_now().await;
4344            }
4345        })
4346        .await
4347        .expect("timed out waiting for inbound WebSocket metadata log");
4348
4349        let (http_task, http_port) = setup_http_test_server().await;
4350        let invalid_headers =
4351            HashMap::from([("x-secret-default".to_string(), format!("{SECRET_MARKER}\n"))]);
4352        let invalid_header_error = HttpClient::builder()
4353            .headers(invalid_headers)
4354            .build()
4355            .unwrap_err();
4356        let http_client = HttpClient::builder().build().unwrap();
4357        let params = HashMap::from([("secret".to_string(), vec![SECRET_MARKER.to_string()])]);
4358        let headers = HashMap::from([
4359            (
4360                "X-Secret-Request".to_string(),
4361                format!("{SECRET_MARKER}:first"),
4362            ),
4363            (
4364                "x-secret-request".to_string(),
4365                format!("{SECRET_MARKER}:second"),
4366            ),
4367        ]);
4368        let http_body = format!("{SECRET_MARKER}:http-body").into_bytes();
4369        http_client
4370            .request(
4371                Method::POST,
4372                format!("http://127.0.0.1:{http_port}/logging"),
4373                Some(&params),
4374                Some(headers),
4375                Some(http_body),
4376                None,
4377                None,
4378            )
4379            .await
4380            .unwrap();
4381
4382        // The initial-connect retry warning names the endpoint, which can carry credentials or
4383        // signed query data, so it must report the redaction placeholder rather than the URL.
4384        let rejecting = TcpListener::bind("127.0.0.1:0").await.unwrap();
4385        let rejecting_port = rejecting.local_addr().unwrap().port();
4386
4387        let rejecting_server = task::spawn(async move {
4388            loop {
4389                let (mut stream, _) = rejecting.accept().await.unwrap();
4390                let mut request = [0; 1024];
4391                let _ = stream.read(&mut request).await.unwrap();
4392                stream
4393                    .write_all(b"HTTP/1.1 503 Service Unavailable\r\n\r\n")
4394                    .await
4395                    .unwrap();
4396            }
4397        });
4398
4399        let retry_config = WebSocketConfig {
4400            url: format!("ws://127.0.0.1:{rejecting_port}/?token={SECRET_MARKER}"),
4401            headers: vec![],
4402            heartbeat_interval_secs: None,
4403            heartbeat_payload: None,
4404            connect_timeout_ms: Some(5_000),
4405            reconnect_delay_initial_ms: None,
4406            reconnect_backoff_factor: None,
4407            reconnect_delay_max_ms: None,
4408            reconnect_jitter_ms: None,
4409            reconnect_max_attempts: None,
4410            heartbeat_timeout_secs: None,
4411            idle_timeout_ms: None,
4412            writer_capacity: None,
4413            backend: TransportBackend::Tungstenite,
4414            proxy_url: None,
4415            max_message_size_bytes: None,
4416            max_frame_size_bytes: None,
4417        };
4418
4419        let retry_error = WebSocketClient::builder()
4420            .config(retry_config)
4421            .message_handler(Arc::new(|_| {}))
4422            .initial_connect_retry_policy(InitialConnectRetryPolicy {
4423                max_attempts: NonZeroU32::new(2).unwrap(),
4424                delay_initial: Duration::from_millis(1),
4425                delay_max: Duration::from_millis(1),
4426                backoff_factor: 1.0,
4427                jitter_ms: 0,
4428            })
4429            .connect()
4430            .await
4431            .expect_err("server always rejects the upgrade");
4432
4433        rejecting_server.abort();
4434
4435        assert!(
4436            matches!(retry_error, TransportError::UpgradeRejected(503)),
4437            "expected a 503 upgrade rejection, was: {retry_error:?}"
4438        );
4439
4440        let messages: Vec<_> = capture
4441            .messages()
4442            .into_iter()
4443            .filter(|(level, message)| {
4444                matches!(level, Level::Trace | Level::Warn)
4445                    && (message.starts_with("Sending ")
4446                        || message.starts_with("Received ")
4447                        || message.starts_with("Replaced ")
4448                        || message.starts_with("WebSocket connection attempt "))
4449            })
4450            .map(|(_, message)| message)
4451            .collect();
4452
4453        let invalid_header_message = invalid_header_error.to_string();
4454
4455        // Asserted positively so the blanket secret assertion below cannot pass vacuously by
4456        // capturing no warning at all.
4457        let retry_warning = messages
4458            .iter()
4459            .find(|message| message.starts_with("WebSocket connection attempt 1/2 "))
4460            .expect("initial-connect retry warning was not captured");
4461        assert!(
4462            retry_warning.contains(REDACTED),
4463            "retry warning omitted the redaction placeholder: {retry_warning}"
4464        );
4465
4466        assert!(
4467            messages.iter().all(|message| {
4468                !message.contains(SECRET_MARKER) && !message.contains(&binary_marker)
4469            }),
4470            "network logs exposed the secret marker: {messages:?}"
4471        );
4472        assert!(
4473            !invalid_header_message.contains(SECRET_MARKER),
4474            "invalid header error exposed the secret marker: {invalid_header_message}"
4475        );
4476        assert!(
4477            invalid_header_message.contains("x-secret-default"),
4478            "invalid header error omitted safe header metadata: {invalid_header_message}"
4479        );
4480        assert!(
4481            messages
4482                .iter()
4483                .any(|message| message == "Sending text frame (28 bytes)"),
4484            "text send metadata missing or inaccurate: {messages:?}"
4485        );
4486        assert!(
4487            messages
4488                .iter()
4489                .any(|message| { message == "Sending text frame once: epoch=0 (31 bytes)" }),
4490            "ownership-bound text metadata missing or inaccurate: {messages:?}"
4491        );
4492        assert!(
4493            messages
4494                .iter()
4495                .any(|message| message == "Sending binary frame (29 bytes)"),
4496            "binary send metadata missing or inaccurate: {messages:?}"
4497        );
4498        assert!(
4499            messages
4500                .iter()
4501                .any(|message| message == "Received text frame (28 bytes)"),
4502            "text receive metadata missing or inaccurate: {messages:?}"
4503        );
4504        assert!(
4505            messages
4506                .iter()
4507                .any(|message| message == "Received text frame (31 bytes)"),
4508            "ownership-bound text receive metadata missing or inaccurate: {messages:?}"
4509        );
4510        assert!(
4511            messages
4512                .iter()
4513                .any(|message| message == "Received message <binary> 29 bytes"),
4514            "binary receive metadata missing or inaccurate: {messages:?}"
4515        );
4516        assert!(
4517            messages
4518                .iter()
4519                .any(|message| message == "Received ping frame (27 bytes)"),
4520            "ping receive metadata missing or inaccurate: {messages:?}"
4521        );
4522        // Two caller headers plus the `accept-encoding: gzip` added for buffered requests
4523        assert!(
4524            messages.iter().any(|message| {
4525                message
4526                    == "Sending HTTP request: method=POST extra_headers=3 query_bytes=29 \
4527                        body_bytes=32"
4528            }),
4529            "HTTP request metadata missing or inaccurate: {messages:?}"
4530        );
4531        assert!(
4532            messages
4533                .iter()
4534                .any(|message| message == "Replaced duplicate request header 'x-secret-request'"),
4535            "duplicate header metadata missing: {messages:?}"
4536        );
4537        assert!(
4538            messages.iter().any(|message| {
4539                message.starts_with("Received HTTP response: status=200 OK headers=")
4540                    && message.ends_with(" body_bytes=31")
4541            }),
4542            "HTTP response metadata missing or inaccurate: {messages:?}"
4543        );
4544
4545        client.disconnect().await;
4546        http_task.abort();
4547    }
4548
4549    // A pong belongs to the connection whose ping caused it: while the
4550    // replacement handshake is gated the pong must be dropped rather than
4551    // waited on, and the next pong must ride the new connection.
4552    #[rstest]
4553    #[tokio::test]
4554    async fn test_send_pong_skips_gated_reconnect() {
4555        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4556        let port = listener.local_addr().unwrap().port();
4557        let (second_accepted_tx, second_accepted_rx) = oneshot::channel();
4558        let (handshake_gate_tx, handshake_gate_rx) = oneshot::channel();
4559        let (observed_tx, mut observed_rx) = mpsc::unbounded_channel();
4560
4561        let server_task = task::spawn(async move {
4562            let (first, _) = listener.accept().await.unwrap();
4563            let first_websocket = accept_async(first).await.unwrap();
4564            drop(first_websocket);
4565
4566            let (second, _) = listener.accept().await.unwrap();
4567            second_accepted_tx.send(()).unwrap();
4568            handshake_gate_rx.await.unwrap();
4569            let mut replacement = accept_async(second).await.unwrap();
4570
4571            while let Some(message) = replacement.next().await {
4572                match message.unwrap() {
4573                    WsMessage::Pong(data) => observed_tx.send(data.to_vec()).unwrap(),
4574                    WsMessage::Close(_) => {
4575                        let _ = replacement.close(None).await;
4576                        break;
4577                    }
4578                    _ => {}
4579                }
4580            }
4581        });
4582
4583        let client = setup_reconnecting_client(port).await;
4584
4585        crate::dst::time::timeout(Duration::from_secs(5), second_accepted_rx)
4586            .await
4587            .expect("replacement connection should be accepted")
4588            .unwrap();
4589        wait_for_mode(&client, ConnectionMode::Reconnect).await;
4590
4591        let rejected = client.send_pong(vec![2; 126]).await;
4592        assert!(matches!(rejected, Err(SendError::InvalidInput(_))));
4593
4594        // The gate is still shut, so a pong that waited for the connection to
4595        // become active could not complete here.
4596        crate::dst::time::timeout(
4597            Duration::from_secs(1),
4598            client.send_pong(b"stale-pong".to_vec()),
4599        )
4600        .await
4601        .expect("pong raised during reconnect should not wait for the replacement")
4602        .expect("skipped pong should report success");
4603
4604        handshake_gate_tx.send(()).unwrap();
4605        wait_for_mode(&client, ConnectionMode::Active).await;
4606
4607        let fresh_payload = b"fresh-pong".to_vec();
4608        client.send_pong(fresh_payload.clone()).await.unwrap();
4609        assert_eq!(
4610            crate::dst::time::timeout(Duration::from_secs(5), observed_rx.recv())
4611                .await
4612                .expect("replacement connection should receive the fresh pong"),
4613            Some(fresh_payload)
4614        );
4615
4616        client.disconnect().await;
4617        server_task.await.unwrap();
4618        assert_eq!(
4619            observed_rx.try_recv(),
4620            Err(mpsc::error::TryRecvError::Disconnected),
4621            "replacement connection should receive exactly one pong"
4622        );
4623    }
4624
4625    // Validation must precede the inactive skip: an oversized pong during
4626    // reconnect is a caller error, not a silently skipped frame.
4627    #[rstest]
4628    #[tokio::test]
4629    async fn test_pong_validation_precedes_inactive_skip() {
4630        let server = TestServer::setup().await;
4631        let client = setup_test_client(server.port).await;
4632        client
4633            .connection_mode
4634            .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
4635
4636        let result = client.send_pong(vec![2; 126]).await;
4637        let epoch_result = client.send_pong_on_connection(vec![2; 126], 0).await;
4638
4639        assert!(matches!(result, Err(SendError::InvalidInput(_))));
4640        assert!(matches!(epoch_result, Err(SendError::InvalidInput(_))));
4641    }
4642
4643    // A rejected oversized pong must not disturb the connection: the frame is
4644    // never enqueued, so the next in-range pong still rides the same socket.
4645    #[rstest]
4646    #[tokio::test]
4647    async fn test_pong_payload_limit_preserves_connection() {
4648        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
4649        let port = listener.local_addr().unwrap().port();
4650        let (pong_tx, mut pong_rx) = mpsc::unbounded_channel();
4651
4652        let server_task = task::spawn(async move {
4653            let (stream, _) = listener.accept().await.unwrap();
4654            let mut websocket = accept_async(stream).await.unwrap();
4655
4656            while let Some(message) = websocket.next().await {
4657                match message.unwrap() {
4658                    WsMessage::Pong(data) => pong_tx.send(data.to_vec()).unwrap(),
4659                    WsMessage::Close(_) => {
4660                        let _ = websocket.close(None).await;
4661                        break;
4662                    }
4663                    _ => {}
4664                }
4665            }
4666        });
4667
4668        let client = setup_reconnecting_client(port).await;
4669        let accepted = vec![1; 125];
4670        let follow_up = vec![3; 125];
4671
4672        client.send_pong(accepted.clone()).await.unwrap();
4673        assert_eq!(
4674            crate::dst::time::timeout(Duration::from_secs(5), pong_rx.recv())
4675                .await
4676                .unwrap(),
4677            Some(accepted)
4678        );
4679
4680        let rejected = client.send_pong(vec![2; 126]).await;
4681        assert!(matches!(rejected, Err(SendError::InvalidInput(_))));
4682
4683        client.send_pong(follow_up.clone()).await.unwrap();
4684        assert_eq!(
4685            crate::dst::time::timeout(Duration::from_secs(5), pong_rx.recv())
4686                .await
4687                .unwrap(),
4688            Some(follow_up)
4689        );
4690
4691        client.disconnect().await;
4692        server_task.await.unwrap();
4693    }
4694
4695    #[tokio::test]
4696    async fn test_websocket_basic() {
4697        let server = TestServer::setup().await;
4698        let client = setup_test_client(server.port).await;
4699
4700        assert!(!client.is_disconnected());
4701
4702        client.disconnect().await;
4703        assert!(client.is_disconnected());
4704    }
4705
4706    #[rstest]
4707    #[tokio::test]
4708    async fn test_drop_sets_shared_connection_mode_closed() {
4709        let server = TestServer::setup().await;
4710        let client = setup_test_client(server.port).await;
4711        let connection_mode = client.connection_mode_atomic();
4712
4713        drop(client);
4714
4715        assert_eq!(
4716            ConnectionMode::from_atomic(&connection_mode),
4717            ConnectionMode::Closed
4718        );
4719    }
4720
4721    #[rstest]
4722    #[tokio::test]
4723    async fn test_notify_closed_closes_reconnecting_client() {
4724        let server = TestServer::setup().await;
4725        let client = setup_test_client(server.port).await;
4726        client
4727            .connection_mode
4728            .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
4729
4730        client.notify_closed();
4731
4732        assert!(client.is_closed());
4733    }
4734
4735    #[tokio::test]
4736    async fn test_websocket_heartbeat() {
4737        let server = TestServer::setup().await;
4738        let client = setup_test_client(server.port).await;
4739
4740        // Wait ~3s => server should see multiple "ping"
4741        tokio::time::sleep(std::time::Duration::from_secs(3)).await;
4742
4743        // Cleanup
4744        client.disconnect().await;
4745        assert!(client.is_disconnected());
4746    }
4747
4748    #[rstest]
4749    #[tokio::test]
4750    async fn test_websocket_reconnect_exhausted() {
4751        let config = WebSocketConfig {
4752            url: "ws://127.0.0.1:9997".into(), // <-- No server
4753            headers: vec![],
4754            heartbeat_interval_secs: None,
4755            heartbeat_payload: None,
4756            connect_timeout_ms: None,
4757            reconnect_delay_initial_ms: None,
4758            reconnect_backoff_factor: None,
4759            reconnect_delay_max_ms: None,
4760            reconnect_jitter_ms: None,
4761            reconnect_max_attempts: None,
4762            heartbeat_timeout_secs: None,
4763            idle_timeout_ms: None,
4764            writer_capacity: None,
4765            backend: TransportBackend::Tungstenite,
4766            proxy_url: None,
4767            max_message_size_bytes: None,
4768            max_frame_size_bytes: None,
4769        };
4770
4771        let states = Arc::new(Mutex::new(Vec::new()));
4772        let states_callback = Arc::clone(&states);
4773        let sink = SocketStateSink::new(move |state| {
4774            states_callback.lock().push(state);
4775        });
4776
4777        let res = WebSocketClient::builder()
4778            .config(config)
4779            .message_handler(Arc::new(|_| {}))
4780            .state_sink(sink)
4781            .connect()
4782            .await;
4783        assert!(res.is_err(), "Should fail quickly with no server");
4784        assert_eq!(*states.lock(), Vec::new());
4785    }
4786
4787    #[tokio::test]
4788    async fn test_websocket_forced_close_reconnect() {
4789        let server = TestServer::setup().await;
4790        let client = setup_test_client(server.port).await;
4791
4792        // 1) Send normal message
4793        client.send_text("Hello".into(), None).await.unwrap();
4794
4795        // 2) Trigger forced close from server
4796        client.send_text("close-now".into(), None).await.unwrap();
4797
4798        // 3) Wait a bit => read loop sees close => reconnect
4799        tokio::time::sleep(std::time::Duration::from_secs(1)).await;
4800
4801        // Confirm not disconnected
4802        assert!(!client.is_disconnected());
4803
4804        // Cleanup
4805        client.disconnect().await;
4806        assert!(client.is_disconnected());
4807    }
4808
4809    #[rstest]
4810    #[case::plain(false)]
4811    #[case::epoch(true)]
4812    #[tokio::test]
4813    async fn test_peer_close_preserves_forwarded_messages(#[case] epoch_aware: bool) {
4814        let timeout = Duration::from_secs(5);
4815        let subscription = "subscribe-retained";
4816        let server = TestServer::setup().await;
4817        let config = WebSocketConfig::builder()
4818            .url(format!("ws://127.0.0.1:{}", server.port))
4819            .headers(vec![("test".to_string(), "test".to_string())])
4820            .backend(TransportBackend::Tungstenite)
4821            .reconnect_delay_initial_ms(50)
4822            .reconnect_jitter_ms(0)
4823            .build()
4824            .unwrap();
4825        let (handler, mut raw_rx) = channel_epoch_message_handler();
4826
4827        let ws_client = if epoch_aware {
4828            WebSocketClient::epoch_builder()
4829                .config(config)
4830                .epoch_handler(handler)
4831                .connect()
4832                .await
4833        } else {
4834            WebSocketClient::builder()
4835                .config(config)
4836                .message_handler(Arc::new(move |message| handler(0, message)))
4837                .connect()
4838                .await
4839        }
4840        .unwrap();
4841
4842        let ws_client = Arc::new(ws_client);
4843        let (bytes_tx, mut bytes_rx) = tokio::sync::mpsc::unbounded_channel();
4844        let forwarding_client = Arc::clone(&ws_client);
4845
4846        let forwarding_task = tokio::spawn(async move {
4847            while let Some((epoch, message)) = raw_rx.recv().await {
4848                let data = match message {
4849                    WsMessage::Text(text) => text.as_bytes().to_vec(),
4850                    WsMessage::Binary(data) => data.to_vec(),
4851                    WsMessage::Close(_) => break,
4852                    _ => continue,
4853                };
4854
4855                if data == RECONNECTED.as_bytes() {
4856                    forwarding_client
4857                        .send_text(subscription.into(), None)
4858                        .await
4859                        .unwrap();
4860                }
4861
4862                if bytes_tx.send((epoch, data)).is_err() {
4863                    break;
4864                }
4865            }
4866        });
4867
4868        ws_client
4869            .send_text(subscription.into(), None)
4870            .await
4871            .unwrap();
4872        let initial = tokio::time::timeout(timeout, bytes_rx.recv())
4873            .await
4874            .expect("initial subscription response timed out")
4875            .expect("forwarding task ended before peer close");
4876
4877        ws_client.send_text("close-now".into(), None).await.unwrap();
4878        let reconnected = tokio::time::timeout(timeout, bytes_rx.recv())
4879            .await
4880            .expect("reconnect notification timed out")
4881            .expect("peer close terminated the forwarding task");
4882        let replayed = tokio::time::timeout(timeout, bytes_rx.recv())
4883            .await
4884            .expect("subscription replay response timed out")
4885            .expect("forwarding task ended before subscription replay");
4886        let payload = vec![3, 7, 11];
4887        ws_client.send_bytes(payload.clone(), None).await.unwrap();
4888        let live = tokio::time::timeout(timeout, bytes_rx.recv())
4889            .await
4890            .expect("replacement connection data timed out")
4891            .expect("forwarding task ended before replacement data");
4892
4893        let expected_epoch = u64::from(epoch_aware);
4894        assert_eq!(initial, (0, subscription.as_bytes().to_vec()));
4895        assert_eq!(
4896            reconnected,
4897            (expected_epoch, RECONNECTED.as_bytes().to_vec())
4898        );
4899        assert_eq!(replayed, (expected_epoch, subscription.as_bytes().to_vec()));
4900        assert_eq!(live, (expected_epoch, payload));
4901        assert_eq!(ws_client.connection_epoch(), 1);
4902        assert!(!forwarding_task.is_finished());
4903
4904        ws_client.disconnect().await;
4905        forwarding_task.abort();
4906    }
4907
4908    #[rstest]
4909    #[case(true)]
4910    #[case(false)]
4911    #[tokio::test]
4912    async fn test_peer_close_frame_is_forwarded_to_close_handler(#[case] reconnect: bool) {
4913        let server = TestServer::setup().await;
4914        let (handler, mut handler_rx) = channel_message_handler();
4915        let (state_tx, mut state_rx) = tokio::sync::mpsc::unbounded_channel();
4916
4917        let client = WebSocketClient::builder()
4918            .config(WebSocketConfig {
4919                url: format!("ws://127.0.0.1:{}", server.port),
4920                headers: vec![("test".into(), "test".into())],
4921                heartbeat_interval_secs: None,
4922                heartbeat_payload: None,
4923                connect_timeout_ms: Some(1_000),
4924                reconnect_delay_initial_ms: Some(1_000),
4925                reconnect_backoff_factor: None,
4926                reconnect_delay_max_ms: None,
4927                reconnect_jitter_ms: Some(0),
4928                reconnect_max_attempts: None,
4929                heartbeat_timeout_secs: None,
4930                idle_timeout_ms: None,
4931                writer_capacity: None,
4932                backend: TransportBackend::Tungstenite,
4933                proxy_url: None,
4934                max_message_size_bytes: None,
4935                max_frame_size_bytes: None,
4936            })
4937            .message_handler(Arc::new(|_| {}))
4938            .state_sink(SocketStateSink::new(move |state| {
4939                state_tx.send(state).expect("state channel closed");
4940            }))
4941            .close_handler(Arc::new(move |frame| {
4942                handler(Message::Close(frame));
4943                reconnect
4944            }))
4945            .connect()
4946            .await
4947            .expect("Failed to connect");
4948
4949        let tracker = AuthTracker::new();
4950        let pending_auth = tracker.begin();
4951        client.set_auth_tracker(tracker.clone(), false);
4952
4953        client
4954            .send_text(CODED_CLOSE_TRIGGER.to_string(), None)
4955            .await
4956            .unwrap();
4957
4958        let message = tokio::time::timeout(Duration::from_secs(10), handler_rx.recv())
4959            .await
4960            .expect("close frame was not forwarded")
4961            .expect("handler channel closed");
4962
4963        let WsMessage::Close(Some(frame)) = message else {
4964            panic!("expected a close frame, was {message:?}");
4965        };
4966
4967        let connected = state_rx.recv().await.expect("initial socket state");
4968        let disconnected = tokio::time::timeout(Duration::from_secs(5), state_rx.recv())
4969            .await
4970            .expect("peer close must publish socket loss")
4971            .expect("state channel closed");
4972
4973        assert_eq!(u16::from(frame.code), CODED_CLOSE_CODE);
4974        assert_eq!(frame.reason.as_str(), CODED_CLOSE_REASON);
4975        assert_eq!(connected, SocketState::Connected);
4976        assert_eq!(disconnected, SocketState::Disconnected);
4977
4978        if !reconnect {
4979            let auth_result = tokio::time::timeout(Duration::from_secs(1), pending_auth)
4980                .await
4981                .expect("terminal close should resolve pending authentication")
4982                .expect("authentication sender should report its terminal result");
4983
4984            assert_eq!(auth_result, Err("WebSocket client closed".to_string()));
4985            assert_eq!(tracker.auth_state(), AuthState::Failed);
4986        }
4987
4988        client.disconnect().await;
4989    }
4990
4991    #[rstest]
4992    #[case(true, ConnectionMode::Active)]
4993    #[case(false, ConnectionMode::Closed)]
4994    fn test_peer_close_callback_controls_reconnect(
4995        #[case] reconnect: bool,
4996        #[case] expected_mode: ConnectionMode,
4997    ) {
4998        let mode = AtomicU8::new(ConnectionMode::Active.as_u8());
4999        let fence = ReadSessionFence::new();
5000        let received = Arc::new(Mutex::new(Vec::new()));
5001        let states = Arc::new(Mutex::new(Vec::new()));
5002        let sink_states = Arc::clone(&states);
5003        let sink = SocketStateSink::new(move |state| sink_states.lock().push(state));
5004        let callback_received = Arc::clone(&received);
5005        let callback: CloseHandler = Arc::new(move |frame| {
5006            callback_received.lock().push(frame);
5007            reconnect
5008        });
5009
5010        let frame = CloseFrame::new(CODED_CLOSE_CODE, CODED_CLOSE_REASON);
5011
5012        WebSocketClientInner::dispatch_close_frame(
5013            &mode,
5014            &fence,
5015            Some(&callback),
5016            Some(&sink),
5017            Some(frame.clone()),
5018        );
5019
5020        assert_eq!(*received.lock(), vec![Some(frame)]);
5021        assert_eq!(ConnectionMode::from_atomic(&mode), expected_mode);
5022        assert_eq!(
5023            *states.lock(),
5024            if reconnect {
5025                Vec::new()
5026            } else {
5027                vec![SocketState::Disconnected]
5028            },
5029        );
5030    }
5031
5032    #[rstest]
5033    fn test_peer_close_callback_ignores_retired_reader() {
5034        let mode = AtomicU8::new(ConnectionMode::Active.as_u8());
5035        let fence = ReadSessionFence::new();
5036        fence.invalidate();
5037        let calls = Arc::new(AtomicUsize::new(0));
5038        let callback_calls = Arc::clone(&calls);
5039        let callback: CloseHandler = Arc::new(move |_| {
5040            callback_calls.fetch_add(1, Ordering::SeqCst);
5041            false
5042        });
5043
5044        WebSocketClientInner::dispatch_close_frame(&mode, &fence, Some(&callback), None, None);
5045
5046        assert_eq!(calls.load(Ordering::SeqCst), 0);
5047        assert_eq!(ConnectionMode::from_atomic(&mode), ConnectionMode::Active);
5048    }
5049
5050    // Server-controlled reasons can contain credentials and forged log lines
5051    #[rstest]
5052    #[tokio::test]
5053    async fn test_close_frame_reason_is_redacted_in_logs() {
5054        let server = TestServer::setup().await;
5055        let client = setup_test_client(server.port).await;
5056        let capture = capture_logs_for(&["nautilus_network::websocket::client"]).await;
5057
5058        client
5059            .send_text(HOSTILE_CLOSE_TRIGGER.to_string(), None)
5060            .await
5061            .unwrap();
5062
5063        tokio::time::timeout(Duration::from_secs(5), async {
5064            loop {
5065                if capture.messages().iter().any(|(_, message)| {
5066                    message.starts_with("Received close frame, terminating: code=")
5067                }) {
5068                    break;
5069                }
5070
5071                tokio::task::yield_now().await;
5072            }
5073        })
5074        .await
5075        .expect("timed out waiting for the close frame log");
5076
5077        let close_frame_logs: Vec<String> = capture
5078            .messages()
5079            .iter()
5080            .map(|(_, message)| message.clone())
5081            .filter(|message| message.starts_with("Received close frame, terminating: code="))
5082            .collect();
5083
5084        let expected = "Received close frame, terminating: code=1011, reason=<redacted>";
5085        assert!(
5086            close_frame_logs.iter().any(|message| message == expected),
5087            "close frame reason was not redacted: {close_frame_logs:?}"
5088        );
5089        assert!(
5090            close_frame_logs.iter().all(|message| {
5091                !message.contains('\r') && !message.contains('\n') && !message.contains('\u{1b}')
5092            }),
5093            "close frame log forged line breaks or escape sequences: {close_frame_logs:?}"
5094        );
5095
5096        assert!(
5097            close_frame_logs
5098                .iter()
5099                .all(|message| !message.contains("injected") && !message.contains("forged"))
5100        );
5101
5102        client.disconnect().await;
5103    }
5104
5105    #[rstest]
5106    #[tokio::test]
5107    async fn test_state_sink_reports_initial_loss_and_recovery() {
5108        let server = TestServer::setup().await;
5109
5110        let config = WebSocketConfig {
5111            url: format!("ws://127.0.0.1:{}", server.port),
5112            headers: vec![("test".into(), "test".into())],
5113            heartbeat_interval_secs: None,
5114            heartbeat_payload: None,
5115            connect_timeout_ms: Some(1_000),
5116            reconnect_delay_initial_ms: Some(1),
5117            reconnect_backoff_factor: Some(1.0),
5118            reconnect_delay_max_ms: Some(1),
5119            reconnect_jitter_ms: Some(0),
5120            reconnect_max_attempts: Some(3),
5121            heartbeat_timeout_secs: None,
5122            idle_timeout_ms: None,
5123            writer_capacity: None,
5124            backend: TransportBackend::Tungstenite,
5125            proxy_url: None,
5126            max_message_size_bytes: None,
5127            max_frame_size_bytes: None,
5128        };
5129
5130        let states = Arc::new(Mutex::new(Vec::new()));
5131        let states_callback = Arc::clone(&states);
5132        let sink = SocketStateSink::new(move |state| {
5133            states_callback.lock().push(state);
5134        });
5135
5136        let client = WebSocketClient::builder()
5137            .config(config)
5138            .message_handler(Arc::new(|_| {}))
5139            .state_sink(sink)
5140            .connect()
5141            .await
5142            .unwrap();
5143
5144        assert_eq!(*states.lock(), vec![SocketState::Connected]);
5145
5146        client.send_text("close-now".into(), None).await.unwrap();
5147        wait_until_async(
5148            || {
5149                let states = Arc::clone(&states);
5150                async move { states.lock().len() == 3 }
5151            },
5152            Duration::from_secs(5),
5153        )
5154        .await;
5155
5156        assert_eq!(
5157            *states.lock(),
5158            vec![
5159                SocketState::Connected,
5160                SocketState::Disconnected,
5161                SocketState::Connected,
5162            ]
5163        );
5164
5165        client.disconnect().await;
5166        assert_eq!(states.lock().len(), 3);
5167    }
5168
5169    #[rstest]
5170    #[tokio::test]
5171    async fn test_drop_suppresses_socket_state_event() {
5172        let server = TestServer::setup().await;
5173
5174        let config = WebSocketConfig {
5175            url: format!("ws://127.0.0.1:{}", server.port),
5176            headers: vec![("test".into(), "test".into())],
5177            heartbeat_interval_secs: None,
5178            heartbeat_payload: None,
5179            connect_timeout_ms: Some(1_000),
5180            reconnect_delay_initial_ms: Some(1),
5181            reconnect_backoff_factor: Some(1.0),
5182            reconnect_delay_max_ms: Some(1),
5183            reconnect_jitter_ms: Some(0),
5184            reconnect_max_attempts: Some(3),
5185            heartbeat_timeout_secs: None,
5186            idle_timeout_ms: None,
5187            writer_capacity: None,
5188            backend: TransportBackend::Tungstenite,
5189            proxy_url: None,
5190            max_message_size_bytes: None,
5191            max_frame_size_bytes: None,
5192        };
5193
5194        let states = Arc::new(Mutex::new(Vec::new()));
5195        let states_callback = Arc::clone(&states);
5196        let sink = SocketStateSink::new(move |state| {
5197            states_callback.lock().push(state);
5198        });
5199
5200        let client = WebSocketClient::builder()
5201            .config(config)
5202            .message_handler(Arc::new(|_| {}))
5203            .state_sink(sink)
5204            .connect()
5205            .await
5206            .unwrap();
5207
5208        drop(client);
5209        crate::dst::time::sleep(Duration::from_millis(25)).await;
5210
5211        assert_eq!(*states.lock(), vec![SocketState::Connected]);
5212    }
5213
5214    #[rstest]
5215    #[tokio::test]
5216    async fn test_stream_state_sink_reports_reader_loss() {
5217        let server = TestServer::setup().await;
5218
5219        let config = WebSocketConfig {
5220            url: format!("ws://127.0.0.1:{}", server.port),
5221            headers: vec![("test".into(), "test".into())],
5222            heartbeat_interval_secs: None,
5223            heartbeat_payload: None,
5224            connect_timeout_ms: None,
5225            reconnect_delay_initial_ms: None,
5226            reconnect_backoff_factor: None,
5227            reconnect_delay_max_ms: None,
5228            reconnect_jitter_ms: None,
5229            reconnect_max_attempts: None,
5230            heartbeat_timeout_secs: None,
5231            idle_timeout_ms: None,
5232            writer_capacity: None,
5233            backend: TransportBackend::Tungstenite,
5234            proxy_url: None,
5235            max_message_size_bytes: None,
5236            max_frame_size_bytes: None,
5237        };
5238
5239        let states = Arc::new(Mutex::new(Vec::new()));
5240        let states_callback = Arc::clone(&states);
5241        let sink = SocketStateSink::new(move |state| {
5242            states_callback.lock().push(state);
5243        });
5244
5245        let (_reader, client) = WebSocketClient::stream_builder()
5246            .config(config)
5247            .state_sink(sink)
5248            .connect()
5249            .await
5250            .unwrap();
5251
5252        client.notify_closed();
5253
5254        assert!(client.is_closed());
5255        assert_eq!(
5256            *states.lock(),
5257            vec![SocketState::Connected, SocketState::Disconnected]
5258        );
5259    }
5260
5261    #[rstest]
5262    #[tokio::test]
5263    async fn test_state_sink_emits_no_retry_or_exhaustion_events() {
5264        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5265        let port = listener.local_addr().unwrap().port();
5266        let server_task = task::spawn(async move {
5267            let (connection, _) = listener.accept().await.unwrap();
5268            let mut websocket = accept_async(connection).await.unwrap();
5269            while let Some(Ok(message)) = websocket.next().await {
5270                if matches!(&message, WsMessage::Text(text) if text.as_str() == "close-now") {
5271                    websocket.close(None).await.unwrap();
5272                    break;
5273                }
5274            }
5275        });
5276
5277        let config = WebSocketConfig {
5278            url: format!("ws://127.0.0.1:{port}"),
5279            headers: vec![],
5280            heartbeat_interval_secs: None,
5281            heartbeat_payload: None,
5282            connect_timeout_ms: Some(100),
5283            reconnect_delay_initial_ms: Some(1),
5284            reconnect_backoff_factor: Some(1.0),
5285            reconnect_delay_max_ms: Some(1),
5286            reconnect_jitter_ms: Some(0),
5287            reconnect_max_attempts: Some(2),
5288            heartbeat_timeout_secs: None,
5289            idle_timeout_ms: None,
5290            writer_capacity: None,
5291            backend: TransportBackend::Tungstenite,
5292            proxy_url: None,
5293            max_message_size_bytes: None,
5294            max_frame_size_bytes: None,
5295        };
5296
5297        let states = Arc::new(Mutex::new(Vec::new()));
5298        let states_callback = Arc::clone(&states);
5299        let sink = SocketStateSink::new(move |state| {
5300            states_callback.lock().push(state);
5301        });
5302
5303        let client = WebSocketClient::builder()
5304            .config(config)
5305            .message_handler(Arc::new(|_| {}))
5306            .state_sink(sink)
5307            .connect()
5308            .await
5309            .unwrap();
5310
5311        client.send_text("close-now".into(), None).await.unwrap();
5312        wait_until_async(
5313            || async { client.is_disconnected() },
5314            Duration::from_secs(5),
5315        )
5316        .await;
5317        assert_eq!(
5318            *states.lock(),
5319            vec![SocketState::Connected, SocketState::Disconnected]
5320        );
5321
5322        server_task.await.unwrap();
5323    }
5324
5325    #[rstest]
5326    #[case(
5327        "invalid header",
5328        "value",
5329        "Invalid WebSocket reconnect header name: invalid HTTP header name"
5330    )]
5331    #[case(
5332        "x-test",
5333        "invalid\nvalue",
5334        "Invalid WebSocket reconnect header value: failed to parse header value"
5335    )]
5336    fn test_reconnect_headers_reject_invalid_update(
5337        #[case] name: &str,
5338        #[case] value: &str,
5339        #[case] expected_message: &str,
5340    ) {
5341        let initial = vec![("x-existing".to_string(), "initial".to_string())];
5342        let headers = ReconnectHeaders::new(initial.clone());
5343
5344        let error = headers.update(name, value).unwrap_err();
5345
5346        let TransportError::Io(error) = error else {
5347            panic!("expected I/O error, was {error:?}");
5348        };
5349
5350        assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
5351        assert_eq!(error.to_string(), expected_message);
5352        assert_eq!(headers.snapshot(), initial);
5353    }
5354
5355    #[tokio::test]
5356    #[allow(clippy::result_large_err)]
5357    async fn test_reconnect_uses_updated_headers_without_interrupting_active_connection() {
5358        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5359        let port = listener.local_addr().unwrap().port();
5360        let (header_tx, mut header_rx) = tokio::sync::mpsc::unbounded_channel();
5361
5362        let server_task = task::spawn(async move {
5363            loop {
5364                let (conn, _) = listener.accept().await.unwrap();
5365                let header_tx = header_tx.clone();
5366
5367                let mut websocket = accept_hdr_async(
5368                    conn,
5369                    move |request: &server::Request, response: server::Response| {
5370                        let values = request
5371                            .headers()
5372                            .get_all("authorization")
5373                            .iter()
5374                            .map(|value| value.to_str().unwrap().to_string())
5375                            .collect::<Vec<_>>();
5376                        header_tx.send(values).unwrap();
5377                        Ok(response)
5378                    },
5379                )
5380                .await
5381                .unwrap();
5382
5383                task::spawn(async move {
5384                    while let Some(Ok(msg)) = websocket.next().await {
5385                        if matches!(&msg, WsMessage::Text(text) if text.as_str() == "close-now") {
5386                            let _ = websocket.close(None).await;
5387                            break;
5388                        }
5389                    }
5390                });
5391            }
5392        });
5393
5394        let config = WebSocketConfig {
5395            url: format!("ws://127.0.0.1:{port}"),
5396            headers: vec![("Authorization".into(), "Bearer initial".into())],
5397            heartbeat_interval_secs: None,
5398            heartbeat_payload: None,
5399            connect_timeout_ms: Some(1_000),
5400            reconnect_delay_initial_ms: Some(50),
5401            reconnect_delay_max_ms: Some(50),
5402            reconnect_backoff_factor: Some(1.0),
5403            reconnect_jitter_ms: Some(0),
5404            reconnect_max_attempts: None,
5405            heartbeat_timeout_secs: None,
5406            idle_timeout_ms: None,
5407            writer_capacity: None,
5408            backend: TransportBackend::Tungstenite,
5409            proxy_url: None,
5410            max_message_size_bytes: None,
5411            max_frame_size_bytes: None,
5412        };
5413
5414        let client = WebSocketClient::builder()
5415            .config(config)
5416            .message_handler(Arc::new(|_| {}))
5417            .connect()
5418            .await
5419            .unwrap();
5420
5421        let initial = header_rx.recv().await.unwrap();
5422        let reconnect_headers = client.reconnect_headers();
5423        reconnect_headers
5424            .update("authorization", "Bearer refreshed")
5425            .unwrap();
5426
5427        tokio::time::sleep(Duration::from_millis(100)).await;
5428        assert_eq!(initial, vec!["Bearer initial"]);
5429        assert!(client.is_active());
5430        assert!(header_rx.try_recv().is_err());
5431        assert!(!format!("{reconnect_headers:?}").contains("refreshed"));
5432
5433        client.send_text("close-now".into(), None).await.unwrap();
5434        let refreshed = tokio::time::timeout(Duration::from_secs(3), header_rx.recv())
5435            .await
5436            .unwrap()
5437            .unwrap();
5438
5439        assert_eq!(refreshed, vec!["Bearer refreshed"]);
5440
5441        client.disconnect().await;
5442        server_task.abort();
5443    }
5444
5445    #[tokio::test]
5446    async fn test_rate_limiter() {
5447        let server = TestServer::setup().await;
5448        let quota = Quota::per_second(NonZeroU32::new(2).unwrap()).unwrap();
5449
5450        let config = WebSocketConfig {
5451            url: format!("ws://127.0.0.1:{}", server.port),
5452            headers: vec![("test".into(), "test".into())],
5453            heartbeat_interval_secs: None,
5454            heartbeat_payload: None,
5455            connect_timeout_ms: None,
5456            reconnect_delay_initial_ms: None,
5457            reconnect_backoff_factor: None,
5458            reconnect_delay_max_ms: None,
5459            reconnect_jitter_ms: None,
5460            reconnect_max_attempts: None,
5461            heartbeat_timeout_secs: None,
5462            idle_timeout_ms: None,
5463            writer_capacity: None,
5464            backend: TransportBackend::Tungstenite,
5465            proxy_url: None,
5466            max_message_size_bytes: None,
5467            max_frame_size_bytes: None,
5468        };
5469
5470        let client = WebSocketClient::builder()
5471            .config(config)
5472            .message_handler(Arc::new(|_| {}))
5473            .keyed_quotas(vec![("default".into(), quota)])
5474            .connect()
5475            .await
5476            .unwrap();
5477
5478        // Burst of 2 passes immediately; the third send must wait for the
5479        // ~500ms replenish interval (keys=None would bypass the limiter)
5480        let keys: [ustr::Ustr; 1] = [ustr::Ustr::from("default")];
5481        let start = std::time::Instant::now();
5482        client
5483            .send_text("test1".into(), Some(keys.as_slice()))
5484            .await
5485            .unwrap();
5486        client
5487            .send_text("test2".into(), Some(keys.as_slice()))
5488            .await
5489            .unwrap();
5490        let after_burst = start.elapsed();
5491        client
5492            .send_text("test3".into(), Some(keys.as_slice()))
5493            .await
5494            .unwrap();
5495        let after_third = start.elapsed();
5496
5497        assert!(
5498            after_burst < std::time::Duration::from_millis(300),
5499            "Burst sends should not be rate limited, took {after_burst:?}"
5500        );
5501        assert!(
5502            after_third >= std::time::Duration::from_millis(400),
5503            "Third send should wait for quota replenishment, took {after_third:?}"
5504        );
5505
5506        // Cleanup
5507        client.disconnect().await;
5508        assert!(client.is_disconnected());
5509    }
5510
5511    #[tokio::test]
5512    async fn test_concurrent_writers() {
5513        let server = TestServer::setup().await;
5514        let client = Arc::new(setup_test_client(server.port).await);
5515
5516        let mut handles = vec![];
5517
5518        for i in 0..10 {
5519            let client = Arc::clone(&client);
5520            handles.push(task::spawn(async move {
5521                client.send_text(format!("test{i}"), None).await.unwrap();
5522            }));
5523        }
5524
5525        for handle in handles {
5526            handle.await.unwrap();
5527        }
5528
5529        // Cleanup
5530        client.disconnect().await;
5531        assert!(client.is_disconnected());
5532    }
5533}
5534
5535#[cfg(test)]
5536#[cfg(not(feature = "turmoil"))]
5537#[cfg(not(all(feature = "simulation", madsim)))] // These tests drive real localhost sockets
5538mod rust_tests {
5539    use std::{
5540        pin::Pin,
5541        sync::{
5542            Arc, OnceLock,
5543            atomic::{AtomicBool, AtomicU8, AtomicUsize, Ordering},
5544        },
5545        task::{Context, Poll},
5546    };
5547
5548    use futures_util::{SinkExt, StreamExt};
5549    #[cfg(target_os = "linux")]
5550    use log::Level;
5551    use nautilus_common::testing::wait_until_async;
5552    use parking_lot::{Condvar, Mutex};
5553    use rstest::rstest;
5554    #[cfg(feature = "transport-sockudo")]
5555    use sockudo_ws::handshake as sockudo_handshake;
5556    #[cfg(feature = "transport-sockudo")]
5557    use tokio::io::AsyncRead;
5558    use tokio::{
5559        io::{AsyncReadExt, AsyncWriteExt},
5560        net::TcpListener,
5561        sync::oneshot,
5562        task::{self, JoinHandle},
5563        time::{Duration, sleep},
5564    };
5565    use tokio_tungstenite::{accept_async, tungstenite::Message as WsMessage};
5566    #[cfg(feature = "transport-sockudo")]
5567    use tokio_tungstenite::{
5568        accept_hdr_async,
5569        tungstenite::{
5570            handshake::server::{self, Callback},
5571            http::HeaderValue,
5572        },
5573    };
5574
5575    use super::*;
5576    #[cfg(target_os = "linux")]
5577    use crate::logging::tests::capture_logs_for;
5578    use crate::{
5579        SocketState,
5580        websocket::types::{channel_epoch_message_handler, channel_message_handler},
5581    };
5582
5583    const TEST_TIMEOUT: Duration = Duration::from_secs(10);
5584
5585    // A retryable upgrade rejection must stay below ERROR on both the transport that reports it
5586    // and the client that retries it.
5587    #[cfg(all(feature = "transport-sockudo", target_os = "linux"))]
5588    const RECONNECT_LOG_TARGETS: &[&str] = &[
5589        "nautilus_network::transport::sockudo",
5590        "nautilus_network::websocket::client",
5591    ];
5592
5593    struct CondvarReleaseGuard<'a> {
5594        release: &'a (Mutex<bool>, Condvar),
5595    }
5596
5597    impl<'a> CondvarReleaseGuard<'a> {
5598        fn new(release: &'a (Mutex<bool>, Condvar)) -> Self {
5599            Self { release }
5600        }
5601
5602        fn release(&self) {
5603            let (lock, condvar) = self.release;
5604            let mut released = lock.lock();
5605            *released = true;
5606            condvar.notify_all();
5607        }
5608    }
5609
5610    impl Drop for CondvarReleaseGuard<'_> {
5611        fn drop(&mut self) {
5612            self.release();
5613        }
5614    }
5615
5616    async fn recv_rendezvous<T: Send + 'static>(
5617        receiver: std::sync::mpsc::Receiver<T>,
5618        name: &'static str,
5619    ) -> T {
5620        let receive_task = tokio::task::spawn_blocking(move || receiver.recv_timeout(TEST_TIMEOUT));
5621
5622        match tokio::time::timeout(TEST_TIMEOUT * 2, receive_task).await {
5623            Ok(Ok(Ok(value))) => value,
5624            Ok(Ok(Err(e))) => {
5625                panic!("{name} did not arrive within the test timeout: {e}")
5626            }
5627            Ok(Err(e)) => panic!("{name} receive task failed: {e}"),
5628            Err(e) => panic!("{name} receive task did not finish: {e}"),
5629        }
5630    }
5631
5632    async fn await_task_termination(task: tokio::task::JoinHandle<()>, name: &'static str) {
5633        match tokio::time::timeout(TEST_TIMEOUT, task).await {
5634            Ok(Ok(())) => {}
5635            Ok(Err(e)) if e.is_cancelled() => {}
5636            Ok(Err(e)) => panic!("{name} failed: {e}"),
5637            Err(e) => panic!("{name} did not terminate within the test timeout: {e}"),
5638        }
5639    }
5640
5641    fn reconnect_test_config(port: u16) -> WebSocketConfig {
5642        WebSocketConfig {
5643            url: format!("ws://127.0.0.1:{port}"),
5644            headers: vec![],
5645            heartbeat_interval_secs: None,
5646            heartbeat_payload: None,
5647            connect_timeout_ms: Some(1_000),
5648            reconnect_delay_initial_ms: None,
5649            reconnect_delay_max_ms: None,
5650            reconnect_backoff_factor: None,
5651            reconnect_jitter_ms: None,
5652            reconnect_max_attempts: None,
5653            heartbeat_timeout_secs: None,
5654            idle_timeout_ms: None,
5655            writer_capacity: None,
5656            backend: TransportBackend::Tungstenite,
5657            proxy_url: None,
5658            max_message_size_bytes: None,
5659            max_frame_size_bytes: None,
5660        }
5661    }
5662
5663    fn initial_connect_retry_policy(
5664        max_attempts: u32,
5665        initial_delay_ms: u64,
5666    ) -> InitialConnectRetryPolicy {
5667        InitialConnectRetryPolicy {
5668            max_attempts: std::num::NonZeroU32::new(max_attempts).unwrap(),
5669            delay_initial: Duration::from_millis(initial_delay_ms),
5670            delay_max: Duration::from_millis(initial_delay_ms),
5671            backoff_factor: 2.0,
5672            jitter_ms: 0,
5673        }
5674    }
5675
5676    #[rstest]
5677    #[case(TransportError::UpgradeRejected(408), true)]
5678    #[case(TransportError::UpgradeRejected(425), true)]
5679    #[case(TransportError::UpgradeRejected(429), true)]
5680    #[case(TransportError::UpgradeRejected(500), true)]
5681    #[case(TransportError::UpgradeRejected(599), true)]
5682    #[case(TransportError::UpgradeRejected(404), false)]
5683    #[case(TransportError::UpgradeRejected(600), false)]
5684    #[case(TransportError::ProxyConnectRejected(429), true)]
5685    #[case(TransportError::ProxyConnectRejected(407), false)]
5686    #[case(TransportError::ConnectionClosed, true)]
5687    #[case(
5688        TransportError::Io(std::io::Error::from(std::io::ErrorKind::TimedOut)),
5689        true
5690    )]
5691    #[case(
5692        TransportError::Io(std::io::Error::from(std::io::ErrorKind::Interrupted)),
5693        true
5694    )]
5695    #[case(
5696        TransportError::Io(std::io::Error::from(std::io::ErrorKind::InvalidInput)),
5697        false
5698    )]
5699    #[case(
5700        TransportError::Io(std::io::Error::from(std::io::ErrorKind::InvalidData)),
5701        false
5702    )]
5703    #[case(
5704        TransportError::Io(std::io::Error::from(std::io::ErrorKind::Unsupported)),
5705        false
5706    )]
5707    #[case(
5708        TransportError::Io(std::io::Error::from(std::io::ErrorKind::PermissionDenied)),
5709        false
5710    )]
5711    #[case(TransportError::Handshake("malformed".to_string()), false)]
5712    fn initial_connect_error_classification(#[case] error: TransportError, #[case] expected: bool) {
5713        assert_eq!(is_retryable_initial_connect_error(&error), expected);
5714    }
5715
5716    async fn rejected_upgrade_attempts(
5717        status: u16,
5718        backend: TransportBackend,
5719        retry_policy: Option<InitialConnectRetryPolicy>,
5720    ) -> u64 {
5721        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5722        let port = listener.local_addr().unwrap().port();
5723        let accepted = Arc::new(AtomicU64::new(0));
5724        let accepted_server = Arc::clone(&accepted);
5725
5726        let server = tokio::spawn(async move {
5727            loop {
5728                let (mut stream, _) = listener.accept().await.unwrap();
5729                accepted_server.fetch_add(1, Ordering::SeqCst);
5730                let mut request = Vec::new();
5731
5732                loop {
5733                    let mut chunk = [0; 1024];
5734                    let n = stream.read(&mut chunk).await.unwrap();
5735                    if n == 0 {
5736                        break;
5737                    }
5738
5739                    request.extend_from_slice(&chunk[..n]);
5740                    if request.windows(4).any(|window| window == b"\r\n\r\n") {
5741                        break;
5742                    }
5743                }
5744
5745                stream
5746                    .write_all(format!("HTTP/1.1 {status}\r\n\r\n").as_bytes())
5747                    .await
5748                    .unwrap();
5749            }
5750        });
5751
5752        let (handler, _rx) = channel_message_handler();
5753        let mut config = reconnect_test_config(port);
5754        config.backend = backend;
5755        WebSocketClient::builder()
5756            .config(config)
5757            .message_handler(handler)
5758            .maybe_initial_connect_retry_policy(retry_policy)
5759            .connect()
5760            .await
5761            .expect_err("server always rejects the upgrade");
5762
5763        let attempts = accepted.load(Ordering::SeqCst);
5764        server.abort();
5765        attempts
5766    }
5767
5768    async fn closed_proxy_attempts() -> u64 {
5769        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5770        let proxy_addr = listener.local_addr().unwrap();
5771        let accepted = Arc::new(AtomicU64::new(0));
5772        let accepted_server = Arc::clone(&accepted);
5773
5774        let server = tokio::spawn(async move {
5775            loop {
5776                let (mut stream, _) = listener.accept().await.unwrap();
5777                accepted_server.fetch_add(1, Ordering::SeqCst);
5778                let mut request = [0; 1024];
5779                let _ = stream.read(&mut request).await.unwrap();
5780            }
5781        });
5782
5783        let (handler, _rx) = channel_message_handler();
5784        let mut config = reconnect_test_config(9);
5785        config.proxy_url = Some(format!("http://{proxy_addr}"));
5786        WebSocketClient::builder()
5787            .config(config)
5788            .message_handler(handler)
5789            .initial_connect_retry_policy(initial_connect_retry_policy(3, 1))
5790            .connect()
5791            .await
5792            .expect_err("proxy always closes before its CONNECT response");
5793
5794        let attempts = accepted.load(Ordering::SeqCst);
5795        server.abort();
5796        attempts
5797    }
5798
5799    #[rstest]
5800    #[tokio::test]
5801    async fn initial_connect_cancellation_interrupts_in_flight_handshake() {
5802        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5803        let port = listener.local_addr().unwrap().port();
5804        let (accepted_tx, accepted_rx) = tokio::sync::oneshot::channel();
5805
5806        let server = tokio::spawn(async move {
5807            let (_stream, _) = listener.accept().await.unwrap();
5808            let _ = accepted_tx.send(());
5809            std::future::pending::<()>().await;
5810        });
5811
5812        let token = CancellationToken::new();
5813        let (handler, _rx) = channel_message_handler();
5814        let mut config = reconnect_test_config(port);
5815        config.connect_timeout_ms = Some(10_000);
5816        let connect = WebSocketClient::builder()
5817            .config(config)
5818            .message_handler(handler)
5819            .initial_connect_retry_policy(initial_connect_retry_policy(5, 30_000))
5820            .cancellation_token(token.clone())
5821            .connect();
5822        tokio::pin!(connect);
5823
5824        tokio::select! {
5825            biased;
5826            result = &mut connect => panic!("connect completed before cancellation: {result:?}"),
5827            result = accepted_rx => result.unwrap(),
5828        }
5829        token.cancel();
5830
5831        let err = tokio::time::timeout(Duration::from_millis(250), connect)
5832            .await
5833            .expect("cancellation should interrupt the handshake")
5834            .expect_err("cancelled connect should fail");
5835        assert!(
5836            matches!(err, TransportError::Io(ref e) if e.kind() == std::io::ErrorKind::Interrupted)
5837        );
5838        server.abort();
5839    }
5840
5841    #[rstest]
5842    #[tokio::test]
5843    async fn timed_out_initial_connect_attempt_is_retried() {
5844        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5845        let port = listener.local_addr().unwrap().port();
5846        let accepted = Arc::new(AtomicU64::new(0));
5847        let accepted_server = Arc::clone(&accepted);
5848
5849        // Accept the TCP connection but never complete the upgrade, so every attempt
5850        // reaches the connect timeout rather than failing fast.
5851        let server = tokio::spawn(async move {
5852            let mut held = Vec::new();
5853
5854            loop {
5855                let (stream, _) = listener.accept().await.unwrap();
5856                accepted_server.fetch_add(1, Ordering::SeqCst);
5857                held.push(stream);
5858            }
5859        });
5860
5861        let (handler, _rx) = channel_message_handler();
5862        let mut config = reconnect_test_config(port);
5863        config.connect_timeout_ms = Some(50);
5864        let err = WebSocketClient::builder()
5865            .config(config)
5866            .message_handler(handler)
5867            .initial_connect_retry_policy(initial_connect_retry_policy(3, 1))
5868            .connect()
5869            .await
5870            .expect_err("every attempt times out");
5871
5872        assert!(
5873            matches!(err, TransportError::Io(ref e) if e.kind() == std::io::ErrorKind::TimedOut)
5874        );
5875
5876        // The ladder must exhaust its attempts: a timeout is transient, so returning after
5877        // the first one leaves this at 1. Wait for the count rather than sampling it, since
5878        // the final connection can complete through the listen backlog before the accepting
5879        // task is scheduled to record it.
5880        let counted = tokio::time::timeout(Duration::from_secs(1), async {
5881            while accepted.load(Ordering::SeqCst) < 3 {
5882                tokio::time::sleep(Duration::from_millis(5)).await;
5883            }
5884        })
5885        .await;
5886
5887        assert!(
5888            counted.is_ok(),
5889            "expected 3 connection attempts, observed {}",
5890            accepted.load(Ordering::SeqCst)
5891        );
5892        // Exactly three: the wait above resolves the scheduling race, but the ladder must
5893        // also stop at its configured maximum rather than exceeding it.
5894        assert_eq!(accepted.load(Ordering::SeqCst), 3);
5895        server.abort();
5896    }
5897
5898    #[rstest]
5899    #[case::too_many_requests(429)]
5900    #[case::server_error(503)]
5901    #[tokio::test]
5902    async fn transient_upgrade_rejection_is_retried(#[case] status: u16) {
5903        assert_eq!(
5904            rejected_upgrade_attempts(
5905                status,
5906                TransportBackend::Tungstenite,
5907                Some(initial_connect_retry_policy(3, 1)),
5908            )
5909            .await,
5910            3
5911        );
5912    }
5913
5914    #[rstest]
5915    #[tokio::test]
5916    async fn initial_connect_without_retry_policy_makes_one_attempt() {
5917        assert_eq!(
5918            rejected_upgrade_attempts(503, TransportBackend::Tungstenite, None).await,
5919            1
5920        );
5921    }
5922
5923    #[cfg(feature = "transport-sockudo")]
5924    #[rstest]
5925    #[tokio::test]
5926    async fn sockudo_too_many_requests_upgrade_rejection_is_retried() {
5927        assert_eq!(
5928            rejected_upgrade_attempts(
5929                429,
5930                TransportBackend::Sockudo,
5931                Some(initial_connect_retry_policy(3, 1)),
5932            )
5933            .await,
5934            3
5935        );
5936    }
5937
5938    #[rstest]
5939    #[tokio::test]
5940    async fn permanent_upgrade_rejection_is_not_retried() {
5941        assert_eq!(
5942            rejected_upgrade_attempts(
5943                404,
5944                TransportBackend::Tungstenite,
5945                Some(initial_connect_retry_policy(3, 1)),
5946            )
5947            .await,
5948            1
5949        );
5950    }
5951
5952    #[rstest]
5953    #[tokio::test]
5954    async fn proxy_close_before_connect_response_is_retried() {
5955        assert_eq!(closed_proxy_attempts().await, 3);
5956    }
5957
5958    #[rstest]
5959    #[tokio::test]
5960    async fn initial_connect_cancellation_interrupts_retry_backoff() {
5961        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
5962        let port = listener.local_addr().unwrap().port();
5963        let accepted = Arc::new(AtomicU64::new(0));
5964        let accepted_server = Arc::clone(&accepted);
5965        let (rejected_tx, rejected_rx) = oneshot::channel();
5966        let server = tokio::spawn(async move {
5967            let mut rejected_tx = Some(rejected_tx);
5968
5969            loop {
5970                let (mut stream, _) = listener.accept().await.unwrap();
5971                accepted_server.fetch_add(1, Ordering::SeqCst);
5972                let mut request = [0; 1024];
5973                let _ = stream.read(&mut request).await.unwrap();
5974                stream
5975                    .write_all(b"HTTP/1.1 503 Service Unavailable\r\n\r\n")
5976                    .await
5977                    .unwrap();
5978
5979                if let Some(rejected_tx) = rejected_tx.take() {
5980                    rejected_tx.send(()).unwrap();
5981                }
5982            }
5983        });
5984
5985        let token = CancellationToken::new();
5986        let (handler, _rx) = channel_message_handler();
5987        let connect = WebSocketClient::builder()
5988            .config(reconnect_test_config(port))
5989            .message_handler(handler)
5990            .initial_connect_retry_policy(initial_connect_retry_policy(5, 30_000))
5991            .cancellation_token(token.clone())
5992            .connect();
5993        tokio::pin!(connect);
5994
5995        tokio::select! {
5996            biased;
5997            result = &mut connect => panic!("connect completed before backoff: {result:?}"),
5998            result = rejected_rx => result.unwrap(),
5999        }
6000        assert!(futures_util::poll!(&mut connect).is_pending());
6001        token.cancel();
6002
6003        let error = tokio::time::timeout(Duration::from_millis(250), connect)
6004            .await
6005            .expect("cancellation should interrupt retry backoff")
6006            .expect_err("cancelled connect should fail");
6007
6008        assert!(
6009            matches!(error, TransportError::Io(ref error) if error.kind() == std::io::ErrorKind::Interrupted)
6010        );
6011        assert_eq!(accepted.load(Ordering::SeqCst), 1);
6012        server.abort();
6013    }
6014
6015    #[rstest]
6016    #[tokio::test]
6017    async fn permanent_initial_connect_error_does_not_wait_for_retry_ladder() {
6018        let (handler, _rx) = channel_message_handler();
6019        let mut config = reconnect_test_config(1);
6020        config.url = "not a websocket URL".to_string();
6021        let connect = WebSocketClient::builder()
6022            .config(config)
6023            .message_handler(handler)
6024            .initial_connect_retry_policy(initial_connect_retry_policy(5, 30_000))
6025            .connect();
6026
6027        let err = tokio::time::timeout(Duration::from_millis(250), connect)
6028            .await
6029            .expect("permanent error should not enter retry backoff")
6030            .expect_err("invalid URL should fail");
6031        assert!(matches!(
6032            err,
6033            TransportError::InvalidUrl(_) | TransportError::Handshake(_)
6034        ));
6035    }
6036
6037    #[rstest]
6038    #[tokio::test(start_paused = true)]
6039    async fn connection_rate_limit_gates_initial_connect_and_reconnect() {
6040        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6041        let port = listener.local_addr().unwrap().port();
6042        let server = tokio::spawn(async move {
6043            let (stream, _) = listener.accept().await.unwrap();
6044            let _first = accept_async(stream).await.unwrap();
6045            let (stream, _) = listener.accept().await.unwrap();
6046            let _second = accept_async(stream).await.unwrap();
6047            std::future::pending::<()>().await;
6048        });
6049
6050        let key = Ustr::from("connection-attempt");
6051
6052        let limiter = Arc::new(RateLimiter::new_with_quota(
6053            None,
6054            vec![(key, Quota::with_period(Duration::from_secs(1)).unwrap())],
6055        ));
6056        limiter.await_keys_ready(Some(&[key])).await;
6057
6058        let rate_limit = ConnectionRateLimit {
6059            limiter,
6060            keys: Arc::from([key]),
6061        };
6062
6063        let (handler, _rx) = channel_message_handler();
6064        let mut config = reconnect_test_config(port);
6065        config.connect_timeout_ms = Some(10_000);
6066        let initial = WebSocketClientInner::connect_url_with_handler(
6067            config,
6068            Some(IncomingHandler::Message(handler)),
6069            None,
6070            None,
6071            None,
6072            Some(rate_limit),
6073            InitialConnectOptions::default(),
6074        );
6075        tokio::pin!(initial);
6076        assert!(futures_util::poll!(&mut initial).is_pending());
6077        tokio::time::advance(Duration::from_secs(1)).await;
6078        tokio::task::yield_now().await;
6079
6080        // Paused time auto-advances while the runtime waits on socket I/O, which
6081        // would fire the connect timeout, so run the handshake on the real clock
6082        tokio::time::resume();
6083        let mut inner = initial.await.unwrap();
6084        tokio::time::pause();
6085
6086        inner
6087            .connection_mode
6088            .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
6089        let reconnect = inner.reconnect_with_outcome();
6090        tokio::pin!(reconnect);
6091        assert!(futures_util::poll!(&mut reconnect).is_pending());
6092        tokio::time::advance(Duration::from_secs(1)).await;
6093        tokio::task::yield_now().await;
6094        tokio::time::resume();
6095        assert_eq!(reconnect.await.unwrap(), ReconnectOutcome::Reconnected);
6096
6097        server.abort();
6098    }
6099
6100    #[rstest]
6101    #[tokio::test(start_paused = true)]
6102    async fn initial_connect_cancellation_interrupts_connection_rate_limit_wait() {
6103        let key = Ustr::from("connection-attempt");
6104
6105        let limiter = Arc::new(RateLimiter::new_with_quota(
6106            None,
6107            vec![(key, Quota::with_period(Duration::from_secs(1)).unwrap())],
6108        ));
6109        limiter.await_keys_ready(Some(&[key])).await;
6110
6111        let rate_limit = ConnectionRateLimit {
6112            limiter,
6113            keys: Arc::from([key]),
6114        };
6115
6116        let token = CancellationToken::new();
6117        let (handler, _rx) = channel_message_handler();
6118
6119        let connect = WebSocketClientInner::connect_url_with_handler(
6120            reconnect_test_config(1),
6121            Some(IncomingHandler::Message(handler)),
6122            None,
6123            None,
6124            None,
6125            Some(rate_limit),
6126            InitialConnectOptions {
6127                retry_policy: None,
6128                cancellation_token: Some(token.clone()),
6129            },
6130        );
6131
6132        tokio::pin!(connect);
6133        assert!(futures_util::poll!(&mut connect).is_pending());
6134
6135        token.cancel();
6136        let error = connect.await.expect_err("cancelled connect should fail");
6137
6138        assert!(
6139            matches!(error, TransportError::Io(ref error) if error.kind() == std::io::ErrorKind::Interrupted)
6140        );
6141    }
6142
6143    #[rstest]
6144    #[tokio::test]
6145    async fn test_reconnect_outcome_is_aborted_before_connect() {
6146        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6147        let port = listener.local_addr().unwrap().port();
6148        let server = tokio::spawn(async move {
6149            let (stream, _) = listener.accept().await.unwrap();
6150            let _websocket = accept_async(stream).await.unwrap();
6151            std::future::pending::<()>().await;
6152        });
6153
6154        let (handler, _rx) = channel_message_handler();
6155        let mut inner =
6156            WebSocketClientInner::connect_url(reconnect_test_config(port), Some(handler), None)
6157                .await
6158                .unwrap();
6159        inner
6160            .connection_mode
6161            .store(ConnectionMode::Disconnect.as_u8(), Ordering::SeqCst);
6162
6163        let outcome = inner.reconnect_with_outcome().await.unwrap();
6164
6165        assert_eq!(outcome, ReconnectOutcome::Aborted);
6166        server.abort();
6167    }
6168
6169    #[rstest]
6170    #[tokio::test]
6171    async fn test_stream_reconnect_outcome_is_aborted_and_notifies_closed() {
6172        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6173        let port = listener.local_addr().unwrap().port();
6174        let server = tokio::spawn(async move {
6175            let (stream, _) = listener.accept().await.unwrap();
6176            let _websocket = accept_async(stream).await.unwrap();
6177            std::future::pending::<()>().await;
6178        });
6179
6180        let mut inner = WebSocketClientInner::connect_url(reconnect_test_config(port), None, None)
6181            .await
6182            .unwrap();
6183        let state_notify = Arc::clone(&inner.state_notify);
6184        let mut notified = std::pin::pin!(state_notify.notified());
6185        notified.as_mut().enable();
6186
6187        let outcome = inner.reconnect_with_outcome().await.unwrap();
6188
6189        assert_eq!(outcome, ReconnectOutcome::Aborted);
6190        assert_eq!(
6191            ConnectionMode::from_atomic(&inner.connection_mode),
6192            ConnectionMode::Closed
6193        );
6194        tokio::time::timeout(TEST_TIMEOUT, notified)
6195            .await
6196            .expect("stream close notification was not published");
6197        server.abort();
6198    }
6199
6200    #[rstest]
6201    #[tokio::test]
6202    async fn test_reconnect_outcome_is_reconnected_with_handler() {
6203        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6204        let port = listener.local_addr().unwrap().port();
6205        let server = tokio::spawn(async move {
6206            let (stream, _) = listener.accept().await.unwrap();
6207            let _first = accept_async(stream).await.unwrap();
6208            let (stream, _) = listener.accept().await.unwrap();
6209            let mut second = accept_async(stream).await.unwrap();
6210            second
6211                .send(WsMessage::Text("replacement".into()))
6212                .await
6213                .unwrap();
6214            std::future::pending::<()>().await;
6215        });
6216
6217        let (epoch_handler, mut epoch_rx) = channel_epoch_message_handler();
6218        let mut inner = WebSocketClientInner::connect_url_with_handler(
6219            reconnect_test_config(port),
6220            Some(IncomingHandler::Epoch(epoch_handler)),
6221            None,
6222            None,
6223            None,
6224            None,
6225            InitialConnectOptions::default(),
6226        )
6227        .await
6228        .unwrap();
6229        inner
6230            .connection_mode
6231            .store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
6232
6233        let outcome = inner.reconnect_with_outcome().await.unwrap();
6234
6235        assert_eq!(outcome, ReconnectOutcome::Reconnected);
6236        assert_eq!(
6237            ConnectionMode::from_atomic(&inner.connection_mode),
6238            ConnectionMode::Active
6239        );
6240        let (epoch, message) = tokio::time::timeout(TEST_TIMEOUT, epoch_rx.recv())
6241            .await
6242            .expect("replacement epoch message was not delivered")
6243            .expect("epoch handler channel closed");
6244        assert_eq!(epoch, 1);
6245        assert_eq!(message, WsMessage::Text("replacement".into()));
6246        server.abort();
6247    }
6248
6249    #[rstest]
6250    #[tokio::test]
6251    async fn test_inner_drop_invalidates_read_fence_and_aborts_tasks() {
6252        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6253        let port = listener.local_addr().unwrap().port();
6254        let server = tokio::spawn(async move {
6255            let (stream, _) = listener.accept().await.unwrap();
6256            let _websocket = accept_async(stream).await.unwrap();
6257            std::future::pending::<()>().await;
6258        });
6259
6260        let mut config = reconnect_test_config(port);
6261        config.heartbeat_interval_secs = Some(60);
6262        let (handler, _handler_rx) = channel_message_handler();
6263        let inner = WebSocketClientInner::connect_url(config, Some(handler), None)
6264            .await
6265            .unwrap();
6266        let read_fence = inner
6267            .read_fence
6268            .clone()
6269            .expect("read fence should exist in handler mode");
6270        let read_abort = inner
6271            .read_task
6272            .as_ref()
6273            .expect("read task should be spawned in handler mode")
6274            .abort_handle();
6275        let write_abort = inner.write_task.abort_handle();
6276        let heartbeat_abort = inner
6277            .heartbeat_task
6278            .as_ref()
6279            .expect("heartbeat task should be spawned for a configured heartbeat")
6280            .abort_handle();
6281
6282        assert!(read_fence.is_valid(), "read fence should start valid");
6283        assert!(
6284            !read_abort.is_finished(),
6285            "read task should be running before drop"
6286        );
6287        assert!(
6288            !write_abort.is_finished(),
6289            "write task should be running before drop"
6290        );
6291        assert!(
6292            !heartbeat_abort.is_finished(),
6293            "heartbeat task should be running before drop"
6294        );
6295
6296        drop(inner);
6297        wait_until_async(
6298            || async {
6299                read_abort.is_finished()
6300                    && write_abort.is_finished()
6301                    && heartbeat_abort.is_finished()
6302            },
6303            TEST_TIMEOUT,
6304        )
6305        .await;
6306
6307        assert!(!read_fence.is_valid(), "read fence was not invalidated");
6308        assert!(read_abort.is_finished(), "read task was not aborted");
6309        assert!(write_abort.is_finished(), "write task was not aborted");
6310        assert!(
6311            heartbeat_abort.is_finished(),
6312            "heartbeat task was not aborted"
6313        );
6314        server.abort();
6315    }
6316
6317    struct RecordingServer {
6318        task: JoinHandle<()>,
6319        port: u16,
6320        messages: Arc<tokio::sync::Mutex<Vec<String>>>,
6321        control_frames: Arc<tokio::sync::Mutex<Vec<String>>>,
6322        connections: Arc<AtomicUsize>,
6323    }
6324
6325    #[cfg(feature = "transport-sockudo")]
6326    async fn read_http_request<S>(stream: &mut S) -> Vec<u8>
6327    where
6328        S: AsyncRead + Unpin,
6329    {
6330        let mut buf = Vec::new();
6331        let mut chunk = [0u8; 256];
6332        loop {
6333            let n = stream.read(&mut chunk).await.unwrap();
6334            assert!(n > 0, "HTTP request closed before headers completed");
6335            buf.extend_from_slice(&chunk[..n]);
6336            if buf.windows(4).any(|window| window == b"\r\n\r\n") {
6337                return buf;
6338            }
6339        }
6340    }
6341
6342    #[cfg(feature = "transport-sockudo")]
6343    fn extract_header<'a>(request: &'a str, name: &str) -> Option<&'a str> {
6344        request.lines().find_map(|line| {
6345            let (header_name, header_value) = line.split_once(':')?;
6346            if header_name.eq_ignore_ascii_case(name) {
6347                Some(header_value.trim())
6348            } else {
6349                None
6350            }
6351        })
6352    }
6353
6354    #[cfg(feature = "transport-sockudo")]
6355    #[derive(Debug, Clone)]
6356    struct HeaderAssertCallback {
6357        key: String,
6358        value: HeaderValue,
6359    }
6360
6361    #[cfg(feature = "transport-sockudo")]
6362    impl Callback for HeaderAssertCallback {
6363        #[expect(
6364            clippy::panic_in_result_fn,
6365            reason = "assertion failures should fail the test"
6366        )]
6367        fn on_request(
6368            self,
6369            request: &server::Request,
6370            response: server::Response,
6371        ) -> Result<server::Response, server::ErrorResponse> {
6372            assert_eq!(request.headers().get(&self.key), Some(&self.value));
6373            Ok(response)
6374        }
6375    }
6376
6377    impl RecordingServer {
6378        async fn setup() -> Self {
6379            let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6380            let port = listener.local_addr().unwrap().port();
6381            let messages = Arc::new(tokio::sync::Mutex::new(Vec::new()));
6382            let messages_clone = Arc::clone(&messages);
6383            let control_frames = Arc::new(tokio::sync::Mutex::new(Vec::new()));
6384            let control_frames_clone = Arc::clone(&control_frames);
6385            let connections = Arc::new(AtomicUsize::new(0));
6386            let connections_clone = Arc::clone(&connections);
6387
6388            let task = task::spawn(async move {
6389                loop {
6390                    let (stream, _) = listener.accept().await.unwrap();
6391                    let mut websocket = accept_async(stream).await.unwrap();
6392                    connections_clone.fetch_add(1, Ordering::SeqCst);
6393                    let messages = Arc::clone(&messages_clone);
6394                    let control_frames = Arc::clone(&control_frames_clone);
6395
6396                    task::spawn(async move {
6397                        while let Some(Ok(msg)) = websocket.next().await {
6398                            match msg {
6399                                WsMessage::Text(text) => {
6400                                    messages.lock().await.push(text.to_string());
6401                                }
6402                                WsMessage::Pong(data) => {
6403                                    control_frames
6404                                        .lock()
6405                                        .await
6406                                        .push(format!("pong:{}", String::from_utf8_lossy(&data)));
6407                                }
6408                                WsMessage::Close(_) => {
6409                                    let _ = websocket.close(None).await;
6410                                    break;
6411                                }
6412                                _ => {}
6413                            }
6414                        }
6415                    });
6416                }
6417            });
6418
6419            Self {
6420                task,
6421                port,
6422                messages,
6423                control_frames,
6424                connections,
6425            }
6426        }
6427
6428        async fn messages(&self) -> Vec<String> {
6429            self.messages.lock().await.clone()
6430        }
6431
6432        async fn control_frames(&self) -> Vec<String> {
6433            self.control_frames.lock().await.clone()
6434        }
6435
6436        async fn wait_for_connections(&self, expected: usize) {
6437            wait_until_async(
6438                || async { self.connections.load(Ordering::SeqCst) == expected },
6439                TEST_TIMEOUT,
6440            )
6441            .await;
6442        }
6443    }
6444
6445    impl Drop for RecordingServer {
6446        fn drop(&mut self) {
6447            self.task.abort();
6448        }
6449    }
6450
6451    #[rstest]
6452    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
6453    async fn test_manual_reconnect_waits_for_slow_loss_callback_and_new_auth() {
6454        let server = RecordingServer::setup().await;
6455        let tracker = AuthTracker::new();
6456        let _initial_auth = tracker.begin();
6457        tracker.succeed();
6458        let states = Arc::new(Mutex::new(Vec::new()));
6459        let states_callback = Arc::clone(&states);
6460        let auth_at_loss = Arc::new(Mutex::new(Vec::new()));
6461        let auth_at_loss_callback = Arc::clone(&auth_at_loss);
6462        let tracker_callback = tracker.clone();
6463        let callback_release = Arc::new((Mutex::new(false), Condvar::new()));
6464        let callback_release_guard = CondvarReleaseGuard::new(callback_release.as_ref());
6465        let callback_release_clone = Arc::clone(&callback_release);
6466        let (callback_entered_tx, callback_entered_rx) = std::sync::mpsc::channel();
6467
6468        let sink = SocketStateSink::new(move |state| {
6469            states_callback.lock().push(state);
6470            if state == SocketState::Disconnected {
6471                auth_at_loss_callback
6472                    .lock()
6473                    .push(tracker_callback.auth_state());
6474                callback_entered_tx.send(()).unwrap();
6475                let (lock, condvar) = callback_release_clone.as_ref();
6476                let mut released = lock.lock();
6477                while !*released {
6478                    condvar.wait(&mut released);
6479                }
6480            }
6481        });
6482
6483        let (handler, mut handler_rx) = channel_message_handler();
6484        let client = WebSocketClient::builder()
6485            .config(reconnect_test_config(server.port))
6486            .message_handler(handler)
6487            .state_sink(sink)
6488            .connect()
6489            .await
6490            .unwrap();
6491        client.set_auth_tracker(tracker.clone(), true);
6492        server.wait_for_connections(1).await;
6493
6494        let handle = client.reconnect_handle();
6495        let (request_tx, request_rx) = std::sync::mpsc::channel();
6496        let request_thread = std::thread::spawn(move || {
6497            request_tx.send(handle.request_reconnect()).unwrap();
6498        });
6499
6500        recv_rendezvous(callback_entered_rx, "slow reconnect callback entry").await;
6501        client
6502            .writer_tx
6503            .send(WriterCommand::Send(Message::text("buffered")))
6504            .unwrap();
6505        server.wait_for_connections(2).await;
6506        tokio::time::sleep(Duration::from_millis(250)).await;
6507
6508        assert_eq!(client.connection_mode(), ConnectionMode::Reconnect);
6509        assert!(!client.reconnect_published.load(Ordering::SeqCst));
6510        assert_eq!(tracker.auth_state(), AuthState::Unauthenticated);
6511        assert_eq!(*auth_at_loss.lock(), vec![AuthState::Unauthenticated]);
6512        assert_eq!(
6513            *states.lock(),
6514            vec![SocketState::Connected, SocketState::Disconnected]
6515        );
6516        assert!(server.messages().await.is_empty());
6517
6518        callback_release_guard.release();
6519        assert_eq!(
6520            recv_rendezvous(request_rx, "manual reconnect result").await,
6521            ReconnectRequestOutcome::Accepted
6522        );
6523        request_thread.join().unwrap();
6524        wait_until_async(|| async { client.is_active() }, TEST_TIMEOUT).await;
6525        wait_until_async(
6526            || {
6527                let states = Arc::clone(&states);
6528                async move { states.lock().len() == 3 }
6529            },
6530            TEST_TIMEOUT,
6531        )
6532        .await;
6533
6534        let notification = tokio::time::timeout(TEST_TIMEOUT, handler_rx.recv())
6535            .await
6536            .expect("reconnect notification was not delivered")
6537            .expect("handler channel closed");
6538        assert_eq!(notification, WsMessage::Text(RECONNECTED.into()));
6539        assert!(handler_rx.try_recv().is_err());
6540        tokio::time::sleep(Duration::from_millis(200)).await;
6541        assert!(server.messages().await.is_empty());
6542
6543        let _replacement_auth = tracker.begin();
6544        tracker.succeed();
6545        wait_until_async(
6546            || {
6547                let messages = Arc::clone(&server.messages);
6548                async move { messages.lock().await.len() == 1 }
6549            },
6550            TEST_TIMEOUT,
6551        )
6552        .await;
6553
6554        client.send_text("live".into(), None).await.unwrap();
6555        wait_until_async(
6556            || {
6557                let messages = Arc::clone(&server.messages);
6558                async move { messages.lock().await.len() == 2 }
6559            },
6560            TEST_TIMEOUT,
6561        )
6562        .await;
6563
6564        assert_eq!(server.messages().await, vec!["buffered", "live"]);
6565        assert_eq!(server.connections.load(Ordering::SeqCst), 2);
6566        assert_eq!(
6567            *states.lock(),
6568            vec![
6569                SocketState::Connected,
6570                SocketState::Disconnected,
6571                SocketState::Connected,
6572            ]
6573        );
6574
6575        client.disconnect().await;
6576    }
6577
6578    #[rstest]
6579    #[tokio::test]
6580    async fn test_reconnect_handle_is_closed_after_client_drop() {
6581        let server = RecordingServer::setup().await;
6582        let (handler, _handler_rx) = channel_message_handler();
6583        let client = WebSocketClient::builder()
6584            .config(reconnect_test_config(server.port))
6585            .message_handler(handler)
6586            .connect()
6587            .await
6588            .unwrap();
6589        let tracker = AuthTracker::new();
6590        let pending_auth = tracker.begin();
6591        client.set_auth_tracker(tracker.clone(), true);
6592        let controller_abort = client.controller_task.abort_handle();
6593        let handle = client.reconnect_handle();
6594
6595        drop(client);
6596        wait_until_async(|| async { controller_abort.is_finished() }, TEST_TIMEOUT).await;
6597
6598        assert_eq!(handle.request_reconnect(), ReconnectRequestOutcome::Closed);
6599        assert_eq!(tracker.auth_state(), AuthState::Failed);
6600        assert_eq!(
6601            tokio::time::timeout(TEST_TIMEOUT, pending_auth)
6602                .await
6603                .expect("client drop should resolve pending authentication")
6604                .expect("authentication sender should report its terminal result"),
6605            Err("WebSocket client closed".to_string())
6606        );
6607    }
6608
6609    #[rstest]
6610    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
6611    async fn test_concurrent_drop_closes_accepted_reconnect() {
6612        let server = RecordingServer::setup().await;
6613        let tracker = AuthTracker::new();
6614        let _initial_auth = tracker.begin();
6615        tracker.succeed();
6616        let callback_release = Arc::new((Mutex::new(false), Condvar::new()));
6617        let callback_release_guard = CondvarReleaseGuard::new(callback_release.as_ref());
6618        let callback_release_clone = Arc::clone(&callback_release);
6619        let (callback_entered_tx, callback_entered_rx) = std::sync::mpsc::channel();
6620        let states = Arc::new(Mutex::new(Vec::new()));
6621        let states_callback = Arc::clone(&states);
6622        let sink = SocketStateSink::new(move |state| {
6623            states_callback.lock().push(state);
6624            if state == SocketState::Disconnected {
6625                callback_entered_tx.send(()).unwrap();
6626                let (lock, condvar) = callback_release_clone.as_ref();
6627                let mut released = lock.lock();
6628                while !*released {
6629                    condvar.wait(&mut released);
6630                }
6631            }
6632        });
6633
6634        let (handler, _handler_rx) = channel_message_handler();
6635        let client = WebSocketClient::builder()
6636            .config(reconnect_test_config(server.port))
6637            .message_handler(handler)
6638            .state_sink(sink)
6639            .connect()
6640            .await
6641            .unwrap();
6642        client.set_auth_tracker(tracker.clone(), true);
6643        let controller_abort = client.controller_task.abort_handle();
6644        let connection_mode = Arc::clone(&client.connection_mode);
6645        let handle = client.reconnect_handle();
6646        let surviving_handle = handle.clone();
6647        let (request_tx, request_rx) = std::sync::mpsc::channel();
6648        let request_thread = std::thread::spawn(move || {
6649            request_tx.send(handle.request_reconnect()).unwrap();
6650        });
6651
6652        recv_rendezvous(callback_entered_rx, "concurrent drop callback entry").await;
6653        drop(client);
6654        wait_until_async(|| async { controller_abort.is_finished() }, TEST_TIMEOUT).await;
6655        assert_eq!(
6656            ConnectionMode::from_atomic(&connection_mode),
6657            ConnectionMode::Closed
6658        );
6659        assert_eq!(tracker.auth_state(), AuthState::Failed);
6660
6661        callback_release_guard.release();
6662        assert_eq!(
6663            recv_rendezvous(request_rx, "concurrent drop reconnect result").await,
6664            ReconnectRequestOutcome::Accepted
6665        );
6666        request_thread.join().unwrap();
6667        assert_eq!(
6668            surviving_handle.request_reconnect(),
6669            ReconnectRequestOutcome::Closed
6670        );
6671        assert_eq!(
6672            *states.lock(),
6673            vec![SocketState::Connected, SocketState::Disconnected]
6674        );
6675    }
6676
6677    #[rstest]
6678    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
6679    async fn test_reconnect_callback_can_drop_client() {
6680        let server = RecordingServer::setup().await;
6681        let client_slot = Arc::new(Mutex::new(None::<WebSocketClient>));
6682        let client_slot_callback = Arc::clone(&client_slot);
6683        let states = Arc::new(Mutex::new(Vec::new()));
6684        let states_callback = Arc::clone(&states);
6685        let sink = SocketStateSink::new(move |state| {
6686            states_callback.lock().push(state);
6687            if state == SocketState::Disconnected {
6688                drop(client_slot_callback.lock().take());
6689            }
6690        });
6691
6692        let (handler, _handler_rx) = channel_message_handler();
6693        let client = WebSocketClient::builder()
6694            .config(reconnect_test_config(server.port))
6695            .message_handler(handler)
6696            .state_sink(sink)
6697            .connect()
6698            .await
6699            .unwrap();
6700        let tracker = AuthTracker::new();
6701        let _initial_auth = tracker.begin();
6702        tracker.succeed();
6703        client.set_auth_tracker(tracker.clone(), true);
6704        let controller_abort = client.controller_task.abort_handle();
6705        let connection_mode = Arc::clone(&client.connection_mode);
6706        let handle = client.reconnect_handle();
6707        let surviving_handle = handle.clone();
6708        *client_slot.lock() = Some(client);
6709        let (result_tx, result_rx) = std::sync::mpsc::channel();
6710        std::thread::spawn(move || {
6711            result_tx.send(handle.request_reconnect()).unwrap();
6712        });
6713
6714        assert_eq!(
6715            recv_rendezvous(result_rx, "callback drop reconnect result").await,
6716            ReconnectRequestOutcome::Accepted
6717        );
6718        wait_until_async(|| async { controller_abort.is_finished() }, TEST_TIMEOUT).await;
6719
6720        assert!(client_slot.lock().is_none());
6721        assert_eq!(
6722            ConnectionMode::from_atomic(&connection_mode),
6723            ConnectionMode::Closed
6724        );
6725        assert_eq!(tracker.auth_state(), AuthState::Failed);
6726        assert_eq!(
6727            surviving_handle.request_reconnect(),
6728            ReconnectRequestOutcome::Closed
6729        );
6730        assert_eq!(
6731            *states.lock(),
6732            vec![SocketState::Connected, SocketState::Disconnected]
6733        );
6734    }
6735
6736    #[rstest]
6737    #[tokio::test]
6738    async fn test_reconnect_then_disconnect() {
6739        // Bind an ephemeral port
6740        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6741        let port = listener.local_addr().unwrap().port();
6742
6743        // Server task: accept one ws connection then close it
6744        let server = task::spawn(async move {
6745            let (stream, _) = listener.accept().await.unwrap();
6746            let ws = accept_async(stream).await.unwrap();
6747            drop(ws);
6748            // Keep alive briefly
6749            sleep(Duration::from_secs(1)).await;
6750        });
6751
6752        // Build a channel-based message handler for incoming messages (unused here)
6753        let (handler, _rx) = channel_message_handler();
6754
6755        // Configure client with short reconnect backoff
6756        let config = WebSocketConfig {
6757            url: format!("ws://127.0.0.1:{port}"),
6758            headers: vec![],
6759            heartbeat_interval_secs: None,
6760            heartbeat_payload: None,
6761            connect_timeout_ms: Some(1_000),
6762            reconnect_delay_initial_ms: Some(50),
6763            reconnect_delay_max_ms: Some(100),
6764            reconnect_backoff_factor: Some(1.0),
6765            reconnect_jitter_ms: Some(0),
6766            reconnect_max_attempts: None,
6767            heartbeat_timeout_secs: None,
6768            idle_timeout_ms: None,
6769            writer_capacity: None,
6770            backend: TransportBackend::Tungstenite,
6771            proxy_url: None,
6772            max_message_size_bytes: None,
6773            max_frame_size_bytes: None,
6774        };
6775
6776        // Connect the client
6777        let client = WebSocketClient::builder()
6778            .config(config)
6779            .message_handler(handler)
6780            .connect()
6781            .await
6782            .unwrap();
6783
6784        // Allow server to drop connection and client to detect
6785        sleep(Duration::from_millis(100)).await;
6786        // Now immediately disconnect the client
6787        client.disconnect().await;
6788        assert!(client.is_disconnected());
6789        server.abort();
6790    }
6791
6792    #[rstest]
6793    #[tokio::test]
6794    async fn test_reconnect_state_flips_when_reader_stops() {
6795        // Bind an ephemeral port and accept a single websocket connection which we drop.
6796        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6797        let port = listener.local_addr().unwrap().port();
6798        let server = task::spawn(async move {
6799            if let Ok((stream, _)) = listener.accept().await
6800                && let Ok(ws) = accept_async(stream).await
6801            {
6802                drop(ws);
6803            }
6804
6805            sleep(Duration::from_millis(50)).await;
6806        });
6807
6808        let (handler, _rx) = channel_message_handler();
6809
6810        let config = WebSocketConfig {
6811            url: format!("ws://127.0.0.1:{port}"),
6812            headers: vec![],
6813            heartbeat_interval_secs: None,
6814            heartbeat_payload: None,
6815            connect_timeout_ms: Some(1_000),
6816            reconnect_delay_initial_ms: Some(50),
6817            reconnect_delay_max_ms: Some(100),
6818            reconnect_backoff_factor: Some(1.0),
6819            reconnect_jitter_ms: Some(0),
6820            reconnect_max_attempts: None,
6821            heartbeat_timeout_secs: None,
6822            idle_timeout_ms: None,
6823            writer_capacity: None,
6824            backend: TransportBackend::Tungstenite,
6825            proxy_url: None,
6826            max_message_size_bytes: None,
6827            max_frame_size_bytes: None,
6828        };
6829
6830        let client = WebSocketClient::builder()
6831            .config(config)
6832            .message_handler(handler)
6833            .connect()
6834            .await
6835            .unwrap();
6836
6837        tokio::time::timeout(Duration::from_secs(2), async {
6838            loop {
6839                if client.is_reconnecting() {
6840                    break;
6841                }
6842
6843                tokio::time::sleep(Duration::from_millis(10)).await;
6844            }
6845        })
6846        .await
6847        .expect("client did not enter RECONNECT state");
6848
6849        client.disconnect().await;
6850        server.abort();
6851    }
6852
6853    #[rstest]
6854    #[tokio::test]
6855    async fn test_stream_mode_disables_auto_reconnect() {
6856        // Test that stream-based clients do not retain an internal message handler
6857        // and that reconnect() transitions to CLOSED state for stream mode
6858        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6859        let port = listener.local_addr().unwrap().port();
6860        let server = task::spawn(async move {
6861            if let Ok((stream, _)) = listener.accept().await
6862                && let Ok(_ws) = accept_async(stream).await
6863            {
6864                // Keep connection alive briefly
6865                sleep(Duration::from_millis(100)).await;
6866            }
6867        });
6868
6869        let config = WebSocketConfig {
6870            url: format!("ws://127.0.0.1:{port}"),
6871            headers: vec![],
6872            heartbeat_interval_secs: None,
6873            heartbeat_payload: None,
6874            connect_timeout_ms: Some(1_000),
6875            reconnect_delay_initial_ms: Some(50),
6876            reconnect_delay_max_ms: Some(100),
6877            reconnect_backoff_factor: Some(1.0),
6878            reconnect_jitter_ms: Some(0),
6879            reconnect_max_attempts: None,
6880            heartbeat_timeout_secs: None,
6881            idle_timeout_ms: None,
6882            writer_capacity: None,
6883            backend: TransportBackend::Tungstenite,
6884            proxy_url: None,
6885            max_message_size_bytes: None,
6886            max_frame_size_bytes: None,
6887        };
6888
6889        let (_reader, _client) = WebSocketClient::stream_builder()
6890            .config(config)
6891            .connect()
6892            .await
6893            .unwrap();
6894
6895        // Note: We can't easily test the reconnect behavior from the outside since
6896        // the inner client is private. The key fix is that WebSocketClientInner
6897        // now has no internal handler in stream mode, and reconnect() will
6898        // transition to CLOSED state instead of creating a new reader that gets dropped.
6899        // This is tested implicitly by the fact that stream users won't get stuck
6900        // in an infinite reconnect loop.
6901
6902        server.abort();
6903    }
6904
6905    #[rstest]
6906    #[tokio::test]
6907    async fn test_message_handler_mode_allows_auto_reconnect() {
6908        // Test that regular clients (with message handler) can auto-reconnect
6909        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6910        let port = listener.local_addr().unwrap().port();
6911
6912        let server = task::spawn(async move {
6913            // Accept first connection and close it
6914            if let Ok((stream, _)) = listener.accept().await
6915                && let Ok(ws) = accept_async(stream).await
6916            {
6917                drop(ws);
6918            }
6919
6920            sleep(Duration::from_millis(50)).await;
6921        });
6922
6923        let (handler, _rx) = channel_message_handler();
6924
6925        let config = WebSocketConfig {
6926            url: format!("ws://127.0.0.1:{port}"),
6927            headers: vec![],
6928            heartbeat_interval_secs: None,
6929            heartbeat_payload: None,
6930            connect_timeout_ms: Some(1_000),
6931            reconnect_delay_initial_ms: Some(50),
6932            reconnect_delay_max_ms: Some(100),
6933            reconnect_backoff_factor: Some(1.0),
6934            reconnect_jitter_ms: Some(0),
6935            reconnect_max_attempts: None,
6936            heartbeat_timeout_secs: None,
6937            idle_timeout_ms: None,
6938            writer_capacity: None,
6939            backend: TransportBackend::Tungstenite,
6940            proxy_url: None,
6941            max_message_size_bytes: None,
6942            max_frame_size_bytes: None,
6943        };
6944
6945        let client = WebSocketClient::builder()
6946            .config(config)
6947            .message_handler(handler)
6948            .connect()
6949            .await
6950            .unwrap();
6951
6952        // Wait for the connection to be dropped and reconnection to be attempted
6953        tokio::time::timeout(Duration::from_secs(2), async {
6954            loop {
6955                if client.is_reconnecting() || client.is_closed() {
6956                    break;
6957                }
6958
6959                tokio::time::sleep(Duration::from_millis(10)).await;
6960            }
6961        })
6962        .await
6963        .expect("client should attempt reconnection or close");
6964
6965        // Should either be reconnecting or closed (depending on timing)
6966        // The important thing is it's not staying active forever
6967        assert!(
6968            client.is_reconnecting() || client.is_closed(),
6969            "Client with message handler should attempt reconnection"
6970        );
6971
6972        client.disconnect().await;
6973        server.abort();
6974    }
6975
6976    #[rstest]
6977    #[tokio::test]
6978    async fn test_handler_mode_reconnect_with_new_connection() {
6979        // Test that handler mode successfully reconnects and messages continue flowing
6980        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
6981        let port = listener.local_addr().unwrap().port();
6982
6983        let server = task::spawn(async move {
6984            // First connection - accept and immediately close
6985            if let Ok((stream, _)) = listener.accept().await
6986                && let Ok(ws) = accept_async(stream).await
6987            {
6988                drop(ws);
6989            }
6990
6991            // Small delay to let client detect disconnection
6992            sleep(Duration::from_millis(100)).await;
6993
6994            // Second connection - accept, send a message, then keep alive
6995            if let Ok((stream, _)) = listener.accept().await
6996                && let Ok(mut ws) = accept_async(stream).await
6997            {
6998                use futures_util::SinkExt;
6999                let _ = ws
7000                    .send(WsMessage::Text("reconnected".to_string().into()))
7001                    .await;
7002                sleep(Duration::from_secs(1)).await;
7003            }
7004        });
7005
7006        let (handler, mut rx) = channel_message_handler();
7007
7008        let config = WebSocketConfig {
7009            url: format!("ws://127.0.0.1:{port}"),
7010            headers: vec![],
7011            heartbeat_interval_secs: None,
7012            heartbeat_payload: None,
7013            connect_timeout_ms: Some(2_000),
7014            reconnect_delay_initial_ms: Some(50),
7015            reconnect_delay_max_ms: Some(200),
7016            reconnect_backoff_factor: Some(1.5),
7017            reconnect_jitter_ms: Some(10),
7018            reconnect_max_attempts: None,
7019            heartbeat_timeout_secs: None,
7020            idle_timeout_ms: None,
7021            writer_capacity: None,
7022            backend: TransportBackend::Tungstenite,
7023            proxy_url: None,
7024            max_message_size_bytes: None,
7025            max_frame_size_bytes: None,
7026        };
7027
7028        let client = WebSocketClient::builder()
7029            .config(config)
7030            .message_handler(handler)
7031            .connect()
7032            .await
7033            .unwrap();
7034
7035        // Wait for reconnection to happen and message to arrive
7036        let result = tokio::time::timeout(Duration::from_secs(5), async {
7037            loop {
7038                if let Ok(msg) = rx.try_recv()
7039                    && matches!(msg, WsMessage::Text(ref text) if AsRef::<str>::as_ref(text) == "reconnected")
7040                {
7041                    return true;
7042                }
7043
7044                tokio::time::sleep(Duration::from_millis(10)).await;
7045            }
7046        })
7047        .await;
7048
7049        assert!(
7050            result.is_ok(),
7051            "Should receive message after reconnection within timeout"
7052        );
7053
7054        client.disconnect().await;
7055        server.abort();
7056    }
7057
7058    #[rstest]
7059    #[tokio::test]
7060    async fn test_stream_mode_no_auto_reconnect() {
7061        // Test that stream mode does not automatically reconnect when connection is lost
7062        // The caller owns the reader and is responsible for detecting disconnection
7063        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7064        let port = listener.local_addr().unwrap().port();
7065
7066        let server = task::spawn(async move {
7067            // Accept connection and send one message, then close
7068            if let Ok((stream, _)) = listener.accept().await
7069                && let Ok(mut ws) = accept_async(stream).await
7070            {
7071                use futures_util::SinkExt;
7072                let _ = ws.send(WsMessage::Text("hello".to_string().into())).await;
7073                sleep(Duration::from_millis(50)).await;
7074                // Connection closes when ws is dropped
7075            }
7076        });
7077
7078        let config = WebSocketConfig {
7079            url: format!("ws://127.0.0.1:{port}"),
7080            headers: vec![],
7081            heartbeat_interval_secs: None,
7082            heartbeat_payload: None,
7083            connect_timeout_ms: Some(1_000),
7084            reconnect_delay_initial_ms: Some(50),
7085            reconnect_delay_max_ms: Some(100),
7086            reconnect_backoff_factor: Some(1.0),
7087            reconnect_jitter_ms: Some(0),
7088            reconnect_max_attempts: None,
7089            heartbeat_timeout_secs: None,
7090            idle_timeout_ms: None,
7091            writer_capacity: None,
7092            backend: TransportBackend::Tungstenite,
7093            proxy_url: None,
7094            max_message_size_bytes: None,
7095            max_frame_size_bytes: None,
7096        };
7097
7098        let (mut reader, client) = WebSocketClient::stream_builder()
7099            .config(config)
7100            .connect()
7101            .await
7102            .unwrap();
7103
7104        // Initially active
7105        assert!(client.is_active(), "Client should start as active");
7106
7107        // Read the hello message
7108        let msg = reader.next().await;
7109        assert!(
7110            matches!(&msg, Some(Ok(Message::Text(bytes))) if bytes.as_ref() == b"hello"),
7111            "Should receive initial message"
7112        );
7113
7114        // Read until connection closes (reader will return None or error)
7115        while let Some(msg) = reader.next().await {
7116            if msg.is_err() || matches!(msg, Ok(Message::Close(_))) {
7117                break;
7118            }
7119        }
7120
7121        // Controller cannot detect reader EOF (reader is owned by caller),
7122        // so the client stays ACTIVE until the caller signals.
7123        sleep(Duration::from_millis(200)).await;
7124        assert!(
7125            client.is_active(),
7126            "Stream mode client stays ACTIVE before notify_closed()"
7127        );
7128
7129        // Caller signals EOF via notify_closed()
7130        client.notify_closed();
7131
7132        assert!(
7133            client.is_closed(),
7134            "Stream mode client should be CLOSED after notify_closed()"
7135        );
7136        assert!(
7137            !client.is_reconnecting(),
7138            "Stream mode client should never attempt reconnection"
7139        );
7140
7141        client.disconnect().await;
7142        server.abort();
7143    }
7144
7145    #[rstest]
7146    #[tokio::test]
7147    async fn test_send_timeout_uses_configured_connect_timeout() {
7148        // Test that send operations respect the configured connect_timeout.
7149        // When a client is stuck in RECONNECT longer than the timeout, sends should fail with Timeout.
7150        use nautilus_common::testing::wait_until_async;
7151
7152        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7153        let port = listener.local_addr().unwrap().port();
7154
7155        let server = task::spawn(async move {
7156            // Accept first connection and immediately close it
7157            if let Ok((stream, _)) = listener.accept().await
7158                && let Ok(ws) = accept_async(stream).await
7159            {
7160                drop(ws);
7161            }
7162
7163            // Don't accept second connection - client will be stuck in RECONNECT
7164            sleep(Duration::from_mins(1)).await;
7165        });
7166
7167        let (handler, _rx) = channel_message_handler();
7168
7169        // Configure with SHORT 2s reconnect timeout
7170        let config = WebSocketConfig {
7171            url: format!("ws://127.0.0.1:{port}"),
7172            headers: vec![],
7173            heartbeat_interval_secs: None,
7174            heartbeat_payload: None,
7175            connect_timeout_ms: Some(2_000), // 2s timeout
7176            reconnect_delay_initial_ms: Some(50),
7177            reconnect_delay_max_ms: Some(100),
7178            reconnect_backoff_factor: Some(1.0),
7179            reconnect_jitter_ms: Some(0),
7180            reconnect_max_attempts: None,
7181            heartbeat_timeout_secs: None,
7182            idle_timeout_ms: None,
7183            writer_capacity: None,
7184            backend: TransportBackend::Tungstenite,
7185            proxy_url: None,
7186            max_message_size_bytes: None,
7187            max_frame_size_bytes: None,
7188        };
7189
7190        let client = WebSocketClient::builder()
7191            .config(config)
7192            .message_handler(handler)
7193            .connect()
7194            .await
7195            .unwrap();
7196
7197        // Wait for client to enter RECONNECT state
7198        wait_until_async(
7199            || async { client.is_reconnecting() },
7200            Duration::from_secs(3),
7201        )
7202        .await;
7203
7204        // Attempt send while stuck in RECONNECT - should timeout after 2s (configured timeout)
7205        let start = std::time::Instant::now();
7206        let send_result = client.send_text("test".to_string(), None).await;
7207        let elapsed = start.elapsed();
7208
7209        assert!(
7210            send_result.is_err(),
7211            "Send should fail when client stuck in RECONNECT"
7212        );
7213        assert!(
7214            matches!(send_result, Err(crate::error::SendError::Timeout)),
7215            "Send should return Timeout error, was: {send_result:?}"
7216        );
7217        // Verify timeout respects configured value (2s), but don't check upper bound
7218        // as CI scheduler jitter can cause legitimate delays beyond the timeout
7219        assert!(
7220            elapsed >= Duration::from_millis(1800),
7221            "Send should timeout after at least 2s (configured timeout), took {elapsed:?}"
7222        );
7223
7224        client.disconnect().await;
7225        server.abort();
7226    }
7227
7228    #[rstest]
7229    #[tokio::test]
7230    async fn test_send_waits_during_reconnection() {
7231        // Test that send operations wait for reconnection to complete (up to timeout)
7232        use nautilus_common::testing::wait_until_async;
7233
7234        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7235        let port = listener.local_addr().unwrap().port();
7236
7237        let server = task::spawn(async move {
7238            // First connection - accept and immediately close
7239            if let Ok((stream, _)) = listener.accept().await
7240                && let Ok(ws) = accept_async(stream).await
7241            {
7242                drop(ws);
7243            }
7244
7245            // Wait a bit before accepting second connection
7246            sleep(Duration::from_millis(500)).await;
7247
7248            // Second connection - accept and keep alive
7249            if let Ok((stream, _)) = listener.accept().await
7250                && let Ok(mut ws) = accept_async(stream).await
7251            {
7252                // Echo messages
7253                while let Some(Ok(msg)) = ws.next().await {
7254                    if ws.send(msg).await.is_err() {
7255                        break;
7256                    }
7257                }
7258            }
7259        });
7260
7261        let (handler, _rx) = channel_message_handler();
7262
7263        let config = WebSocketConfig {
7264            url: format!("ws://127.0.0.1:{port}"),
7265            headers: vec![],
7266            heartbeat_interval_secs: None,
7267            heartbeat_payload: None,
7268            connect_timeout_ms: Some(5_000), // 5s timeout - enough for reconnect
7269            reconnect_delay_initial_ms: Some(100),
7270            reconnect_delay_max_ms: Some(200),
7271            reconnect_backoff_factor: Some(1.0),
7272            reconnect_jitter_ms: Some(0),
7273            reconnect_max_attempts: None,
7274            heartbeat_timeout_secs: None,
7275            idle_timeout_ms: None,
7276            writer_capacity: None,
7277            backend: TransportBackend::Tungstenite,
7278            proxy_url: None,
7279            max_message_size_bytes: None,
7280            max_frame_size_bytes: None,
7281        };
7282
7283        let client = WebSocketClient::builder()
7284            .config(config)
7285            .message_handler(handler)
7286            .connect()
7287            .await
7288            .unwrap();
7289
7290        // Wait for reconnection to trigger
7291        wait_until_async(
7292            || async { client.is_reconnecting() },
7293            Duration::from_secs(2),
7294        )
7295        .await;
7296
7297        // Try to send while reconnecting - should wait and succeed after reconnect
7298        let send_result = tokio::time::timeout(
7299            Duration::from_secs(3),
7300            client.send_text("test_message".to_string(), None),
7301        )
7302        .await;
7303
7304        assert!(
7305            send_result.is_ok() && send_result.unwrap().is_ok(),
7306            "Send should succeed after waiting for reconnection"
7307        );
7308
7309        client.disconnect().await;
7310        server.abort();
7311    }
7312
7313    #[rstest]
7314    #[tokio::test]
7315    async fn test_rate_limiter_before_active_wait() {
7316        // Test that rate limiting happens BEFORE active state check.
7317        // This prevents race conditions where connection state changes during rate limit wait.
7318        // We verify this by: (1) exhausting rate limit, (2) ensuring client is RECONNECTING,
7319        // (3) sending again and confirming it waits for rate limit THEN reconnection.
7320        use std::{num::NonZeroU32, sync::Arc};
7321
7322        use nautilus_common::testing::wait_until_async;
7323
7324        use crate::ratelimiter::quota::Quota;
7325
7326        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7327        let port = listener.local_addr().unwrap().port();
7328
7329        let server = task::spawn(async move {
7330            // First connection - accept and close after receiving one message
7331            if let Ok((stream, _)) = listener.accept().await
7332                && let Ok(mut ws) = accept_async(stream).await
7333            {
7334                // Receive first message then close
7335                if let Some(Ok(_)) = ws.next().await {
7336                    drop(ws);
7337                }
7338            }
7339
7340            // Wait before accepting reconnection
7341            sleep(Duration::from_millis(500)).await;
7342
7343            // Second connection - accept and keep alive
7344            if let Ok((stream, _)) = listener.accept().await
7345                && let Ok(mut ws) = accept_async(stream).await
7346            {
7347                while let Some(Ok(msg)) = ws.next().await {
7348                    if ws.send(msg).await.is_err() {
7349                        break;
7350                    }
7351                }
7352            }
7353        });
7354
7355        let (handler, _rx) = channel_message_handler();
7356
7357        let config = WebSocketConfig {
7358            url: format!("ws://127.0.0.1:{port}"),
7359            headers: vec![],
7360            heartbeat_interval_secs: None,
7361            heartbeat_payload: None,
7362            connect_timeout_ms: Some(5_000),
7363            reconnect_delay_initial_ms: Some(50),
7364            reconnect_delay_max_ms: Some(100),
7365            reconnect_backoff_factor: Some(1.0),
7366            reconnect_jitter_ms: Some(0),
7367            reconnect_max_attempts: None,
7368            heartbeat_timeout_secs: None,
7369            idle_timeout_ms: None,
7370            writer_capacity: None,
7371            backend: TransportBackend::Tungstenite,
7372            proxy_url: None,
7373            max_message_size_bytes: None,
7374            max_frame_size_bytes: None,
7375        };
7376
7377        // Very restrictive rate limit: 1 request per second, burst of 1
7378        let quota = Quota::per_second(NonZeroU32::new(1).unwrap())
7379            .unwrap()
7380            .allow_burst(NonZeroU32::new(1).unwrap());
7381
7382        let client = Arc::new(
7383            WebSocketClient::builder()
7384                .config(config)
7385                .message_handler(handler)
7386                .keyed_quotas(vec![("test_key".to_string(), quota)])
7387                .connect()
7388                .await
7389                .unwrap(),
7390        );
7391
7392        // First send exhausts burst capacity and triggers connection close
7393        let test_key: [Ustr; 1] = [Ustr::from("test_key")];
7394        client
7395            .send_text("msg1".to_string(), Some(test_key.as_slice()))
7396            .await
7397            .unwrap();
7398
7399        // Wait for client to enter RECONNECT state
7400        wait_until_async(
7401            || async { client.is_reconnecting() },
7402            Duration::from_secs(2),
7403        )
7404        .await;
7405
7406        // Second send: will hit rate limit (~1s) THEN wait for reconnection (~0.5s)
7407        let start = std::time::Instant::now();
7408        let send_result = client
7409            .send_text("msg2".to_string(), Some(test_key.as_slice()))
7410            .await;
7411        let elapsed = start.elapsed();
7412
7413        // Should succeed after both rate limit AND reconnection
7414        assert!(
7415            send_result.is_ok(),
7416            "Send should succeed after rate limit + reconnection, was: {send_result:?}"
7417        );
7418        // Total wait should be at least rate limit time (~1s)
7419        // The reconnection completes while rate limiting or after
7420        // Use 850ms threshold to account for timing jitter in CI
7421        assert!(
7422            elapsed >= Duration::from_millis(850),
7423            "Should wait for rate limit (~1s), waited {elapsed:?}"
7424        );
7425
7426        client.disconnect().await;
7427        server.abort();
7428    }
7429
7430    #[rstest]
7431    #[tokio::test]
7432    async fn test_disconnect_during_reconnect_exits_cleanly() {
7433        // Test CAS race condition: disconnect called during reconnection
7434        // Should exit cleanly without spawning new tasks
7435        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7436        let port = listener.local_addr().unwrap().port();
7437
7438        let server = task::spawn(async move {
7439            // Accept first connection and immediately close
7440            if let Ok((stream, _)) = listener.accept().await
7441                && let Ok(ws) = accept_async(stream).await
7442            {
7443                drop(ws);
7444            }
7445
7446            // Don't accept second connection - let reconnect hang
7447            sleep(Duration::from_mins(1)).await;
7448        });
7449
7450        let (handler, _rx) = channel_message_handler();
7451
7452        let config = WebSocketConfig {
7453            url: format!("ws://127.0.0.1:{port}"),
7454            headers: vec![],
7455            heartbeat_interval_secs: None,
7456            heartbeat_payload: None,
7457            connect_timeout_ms: Some(2_000), // 2s timeout - shorter than disconnect timeout
7458            reconnect_delay_initial_ms: Some(100),
7459            reconnect_delay_max_ms: Some(200),
7460            reconnect_backoff_factor: Some(1.0),
7461            reconnect_jitter_ms: Some(0),
7462            reconnect_max_attempts: None,
7463            heartbeat_timeout_secs: None,
7464            idle_timeout_ms: None,
7465            writer_capacity: None,
7466            backend: TransportBackend::Tungstenite,
7467            proxy_url: None,
7468            max_message_size_bytes: None,
7469            max_frame_size_bytes: None,
7470        };
7471
7472        let client = WebSocketClient::builder()
7473            .config(config)
7474            .message_handler(handler)
7475            .connect()
7476            .await
7477            .unwrap();
7478
7479        // Wait for reconnection to start
7480        tokio::time::timeout(Duration::from_secs(2), async {
7481            while !client.is_reconnecting() {
7482                sleep(Duration::from_millis(10)).await;
7483            }
7484        })
7485        .await
7486        .expect("Client should enter RECONNECT state");
7487
7488        // Disconnect while reconnecting
7489        client.disconnect().await;
7490
7491        // Should be cleanly closed
7492        assert!(
7493            client.is_disconnected(),
7494            "Client should be cleanly disconnected"
7495        );
7496
7497        server.abort();
7498    }
7499
7500    #[rstest]
7501    #[tokio::test]
7502    async fn test_send_fails_fast_when_closed_before_rate_limit() {
7503        // Test that send operations check connection state BEFORE rate limiting,
7504        // preventing unnecessary delays when the connection is already closed.
7505        use std::{num::NonZeroU32, sync::Arc};
7506
7507        use nautilus_common::testing::wait_until_async;
7508
7509        use crate::ratelimiter::quota::Quota;
7510
7511        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7512        let port = listener.local_addr().unwrap().port();
7513
7514        let server = task::spawn(async move {
7515            // Accept connection and immediately close
7516            if let Ok((stream, _)) = listener.accept().await
7517                && let Ok(ws) = accept_async(stream).await
7518            {
7519                drop(ws);
7520            }
7521
7522            sleep(Duration::from_mins(1)).await;
7523        });
7524
7525        let (handler, _rx) = channel_message_handler();
7526
7527        let config = WebSocketConfig {
7528            url: format!("ws://127.0.0.1:{port}"),
7529            headers: vec![],
7530            heartbeat_interval_secs: None,
7531            heartbeat_payload: None,
7532            connect_timeout_ms: Some(5_000),
7533            reconnect_delay_initial_ms: Some(50),
7534            reconnect_delay_max_ms: Some(100),
7535            reconnect_backoff_factor: Some(1.0),
7536            reconnect_jitter_ms: Some(0),
7537            reconnect_max_attempts: None,
7538            heartbeat_timeout_secs: None,
7539            idle_timeout_ms: None,
7540            writer_capacity: None,
7541            backend: TransportBackend::Tungstenite,
7542            proxy_url: None,
7543            max_message_size_bytes: None,
7544            max_frame_size_bytes: None,
7545        };
7546
7547        // Very restrictive rate limit: 1 request per 10 seconds
7548        // This ensures that if we wait for rate limit, the test will timeout
7549        let quota = Quota::with_period(Duration::from_secs(10))
7550            .unwrap()
7551            .allow_burst(NonZeroU32::new(1).unwrap());
7552
7553        let client = Arc::new(
7554            WebSocketClient::builder()
7555                .config(config)
7556                .message_handler(handler)
7557                .keyed_quotas(vec![("test_key".to_string(), quota)])
7558                .connect()
7559                .await
7560                .unwrap(),
7561        );
7562
7563        // Wait for disconnection
7564        wait_until_async(
7565            || async { client.is_reconnecting() || client.is_closed() },
7566            Duration::from_secs(2),
7567        )
7568        .await;
7569
7570        // Explicitly disconnect to move away from ACTIVE state
7571        client.disconnect().await;
7572        assert!(
7573            !client.is_active(),
7574            "Client should not be active after disconnect"
7575        );
7576
7577        // Attempt send - should fail IMMEDIATELY without waiting for rate limit
7578        let start = std::time::Instant::now();
7579        let test_key: [Ustr; 1] = [Ustr::from("test_key")];
7580        let result = client
7581            .send_text("test".to_string(), Some(test_key.as_slice()))
7582            .await;
7583        let elapsed = start.elapsed();
7584
7585        // Should fail with Closed error
7586        assert!(result.is_err(), "Send should fail when client is closed");
7587        assert!(
7588            matches!(result, Err(crate::error::SendError::Closed)),
7589            "Send should return Closed error, was: {result:?}"
7590        );
7591
7592        // Should fail FAST (< 100ms) without waiting for rate limit (10s)
7593        assert!(
7594            elapsed < Duration::from_millis(100),
7595            "Send should fail fast without rate limiting, took {elapsed:?}"
7596        );
7597
7598        server.abort();
7599    }
7600
7601    #[rstest]
7602    #[tokio::test]
7603    async fn test_builder_rejects_shared_rate_limiter_with_quotas() {
7604        let (handler, _rx) = channel_message_handler();
7605        let rate_limiter = Arc::new(RateLimiter::new_with_quota(None, Vec::new()));
7606        let result = WebSocketClient::builder()
7607            .config(reconnect_test_config(1))
7608            .message_handler(handler)
7609            .default_quota(Quota::with_period(Duration::from_secs(1)).unwrap())
7610            .rate_limiter(rate_limiter)
7611            .connect()
7612            .await;
7613
7614        assert_eq!(
7615            result.unwrap_err().to_string(),
7616            "I/O error: Cannot combine a shared rate limiter with quota configuration"
7617        );
7618    }
7619
7620    #[rstest]
7621    #[tokio::test]
7622    async fn test_builder_rejects_connection_keys_without_limiter() {
7623        let (handler, _rx) = channel_message_handler();
7624        let result = WebSocketClient::builder()
7625            .config(reconnect_test_config(1))
7626            .message_handler(handler)
7627            .connection_rate_keys(Arc::from([Ustr::from("connect")]))
7628            .connect()
7629            .await;
7630
7631        assert_eq!(
7632            result.unwrap_err().to_string(),
7633            "I/O error: Connection rate keys require a connection rate limiter"
7634        );
7635    }
7636
7637    #[rstest]
7638    #[tokio::test]
7639    async fn test_builder_rejects_connection_limiter_without_keys() {
7640        let (handler, _rx) = channel_message_handler();
7641        let rate_limiter = Arc::new(RateLimiter::new_with_quota(None, Vec::new()));
7642        let result = WebSocketClient::builder()
7643            .config(reconnect_test_config(1))
7644            .message_handler(handler)
7645            .connection_rate_limiter(rate_limiter)
7646            .connect()
7647            .await;
7648
7649        assert_eq!(
7650            result.unwrap_err().to_string(),
7651            "I/O error: Connection rate limiter requires at least one connection rate key"
7652        );
7653    }
7654
7655    #[rstest]
7656    #[tokio::test]
7657    async fn test_epoch_builder_rejects_both_ping_handler_types() {
7658        let (handler, _rx) = channel_epoch_message_handler();
7659        let ping_handler: PingHandler = Arc::new(|_| {});
7660        let epoch_ping_handler: EpochPingHandler = Arc::new(|_, _| {});
7661        let result = WebSocketClient::epoch_builder()
7662            .config(reconnect_test_config(1))
7663            .epoch_handler(handler)
7664            .ping_handler(ping_handler)
7665            .epoch_ping_handler(epoch_ping_handler)
7666            .connect()
7667            .await;
7668
7669        assert_eq!(
7670            result.unwrap_err().to_string(),
7671            "I/O error: Cannot configure both ping_handler and epoch_ping_handler"
7672        );
7673    }
7674
7675    #[rstest]
7676    #[tokio::test]
7677    async fn test_connect_url_rejects_invalid_reconnect_timing_before_connect() {
7678        let (handler, _rx) = channel_message_handler();
7679
7680        let config = WebSocketConfig {
7681            url: "ws://127.0.0.1:1".to_string(),
7682            headers: vec![],
7683            heartbeat_interval_secs: None,
7684            heartbeat_payload: None,
7685            connect_timeout_ms: Some(0),
7686            reconnect_delay_initial_ms: Some(100),
7687            reconnect_delay_max_ms: Some(500),
7688            reconnect_backoff_factor: Some(1.5),
7689            reconnect_jitter_ms: Some(0),
7690            reconnect_max_attempts: None,
7691            heartbeat_timeout_secs: None,
7692            idle_timeout_ms: None,
7693            writer_capacity: None,
7694            backend: TransportBackend::Tungstenite,
7695            proxy_url: None,
7696            max_message_size_bytes: None,
7697            max_frame_size_bytes: None,
7698        };
7699
7700        let err = WebSocketClientInner::connect_url(config, Some(handler), None)
7701            .await
7702            .expect_err("invalid reconnect timing should be rejected");
7703
7704        match err {
7705            TransportError::Io(error) => {
7706                assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
7707                assert!(
7708                    error.to_string().contains("connect_timeout_ms"),
7709                    "error should mention zero reconnect timeout, was: {error}"
7710                );
7711            }
7712            other => panic!("expected InvalidInput IO error, was: {other:?}"),
7713        }
7714    }
7715
7716    #[rstest]
7717    #[tokio::test]
7718    async fn test_connect_url_rejects_invalid_reconnect_backoff_before_connect() {
7719        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7720        let port = listener.local_addr().unwrap().port();
7721        let accepted = Arc::new(std::sync::atomic::AtomicBool::new(false));
7722        let accepted_clone = Arc::clone(&accepted);
7723
7724        let server = task::spawn(async move {
7725            let (stream, _) = listener.accept().await.unwrap();
7726            accepted_clone.store(true, Ordering::SeqCst);
7727            accept_async(stream).await.unwrap();
7728        });
7729
7730        let (handler, _rx) = channel_message_handler();
7731
7732        let config = WebSocketConfig {
7733            url: format!("ws://127.0.0.1:{port}"),
7734            headers: vec![],
7735            heartbeat_interval_secs: None,
7736            heartbeat_payload: None,
7737            connect_timeout_ms: Some(1_000),
7738            reconnect_delay_initial_ms: Some(50),
7739            reconnect_delay_max_ms: Some(100),
7740            reconnect_backoff_factor: Some(100.1),
7741            reconnect_jitter_ms: Some(0),
7742            reconnect_max_attempts: None,
7743            heartbeat_timeout_secs: None,
7744            idle_timeout_ms: None,
7745            writer_capacity: None,
7746            backend: TransportBackend::Tungstenite,
7747            proxy_url: None,
7748            max_message_size_bytes: None,
7749            max_frame_size_bytes: None,
7750        };
7751
7752        let error = WebSocketClientInner::connect_url(config, Some(handler), None)
7753            .await
7754            .expect_err("invalid reconnect backoff should be rejected");
7755
7756        match error {
7757            TransportError::Io(error) => {
7758                assert_eq!(error.kind(), std::io::ErrorKind::InvalidInput);
7759                assert!(
7760                    error.to_string().contains("factor"),
7761                    "error should mention the invalid factor, was: {error}"
7762                );
7763            }
7764            other => panic!("expected InvalidInput IO error, was: {other:?}"),
7765        }
7766
7767        assert!(
7768            !accepted.load(Ordering::SeqCst),
7769            "invalid reconnect backoff must be rejected before connecting"
7770        );
7771        server.abort();
7772    }
7773
7774    #[rstest]
7775    #[tokio::test]
7776    async fn test_client_without_handler_sets_stream_mode() {
7777        // Test that if a client is created without a handler via connect_url,
7778        // it keeps the internal handler empty to prevent zombie connections
7779
7780        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7781        let port = listener.local_addr().unwrap().port();
7782
7783        let server = task::spawn(async move {
7784            // Accept and immediately close to simulate server disconnect
7785            if let Ok((stream, _)) = listener.accept().await
7786                && let Ok(ws) = accept_async(stream).await
7787            {
7788                drop(ws); // Drop connection immediately
7789            }
7790        });
7791
7792        let config = WebSocketConfig {
7793            url: format!("ws://127.0.0.1:{port}"),
7794            headers: vec![],
7795            heartbeat_interval_secs: None,
7796            heartbeat_payload: None,
7797            connect_timeout_ms: Some(1_000),
7798            reconnect_delay_initial_ms: Some(100),
7799            reconnect_delay_max_ms: Some(500),
7800            reconnect_backoff_factor: Some(1.5),
7801            reconnect_jitter_ms: Some(0),
7802            reconnect_max_attempts: None,
7803            heartbeat_timeout_secs: None,
7804            idle_timeout_ms: None,
7805            writer_capacity: None,
7806            backend: TransportBackend::Tungstenite,
7807            proxy_url: None,
7808            max_message_size_bytes: None,
7809            max_frame_size_bytes: None,
7810        };
7811
7812        // Create client directly via connect_url with no handler (stream mode)
7813        let inner = WebSocketClientInner::connect_url(config, None, None)
7814            .await
7815            .unwrap();
7816
7817        // Verify stream mode does not retain an internal handler
7818        assert!(
7819            inner.handler.is_none(),
7820            "Client without handler should not retain an internal handler"
7821        );
7822
7823        // Verify that when stream mode is enabled, reconnection is disabled
7824        // (documented behavior - stream mode clients close instead of reconnecting)
7825
7826        server.abort();
7827    }
7828
7829    #[rstest]
7830    #[tokio::test]
7831    async fn test_idle_timeout_triggers_reconnect() {
7832        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7833        let port = listener.local_addr().unwrap().port();
7834
7835        // Server accepts WS connection but sends nothing (simulates silent death)
7836        let server = task::spawn(async move {
7837            let (stream, _) = listener.accept().await.unwrap();
7838            let _ws = accept_async(stream).await.unwrap();
7839            // Hold connection open but send nothing
7840            sleep(Duration::from_secs(5)).await;
7841        });
7842
7843        let (handler, _rx) = channel_message_handler();
7844
7845        let config = WebSocketConfig {
7846            url: format!("ws://127.0.0.1:{port}"),
7847            headers: vec![],
7848            heartbeat_interval_secs: None,
7849            heartbeat_payload: None,
7850            connect_timeout_ms: Some(2_000),
7851            reconnect_delay_initial_ms: Some(50),
7852            reconnect_delay_max_ms: Some(100),
7853            reconnect_backoff_factor: Some(1.0),
7854            reconnect_jitter_ms: Some(0),
7855            reconnect_max_attempts: Some(1),
7856            heartbeat_timeout_secs: None,
7857            idle_timeout_ms: Some(500),
7858            writer_capacity: None,
7859            backend: TransportBackend::Tungstenite,
7860            proxy_url: None,
7861            max_message_size_bytes: None,
7862            max_frame_size_bytes: None,
7863        };
7864
7865        let client = WebSocketClient::builder()
7866            .config(config)
7867            .message_handler(handler)
7868            .connect()
7869            .await
7870            .unwrap();
7871
7872        assert!(client.is_active());
7873
7874        // Wait for idle timeout to fire and client to enter reconnect/closed
7875        wait_until_async(
7876            || async { client.is_reconnecting() || client.is_disconnected() },
7877            Duration::from_secs(3),
7878        )
7879        .await;
7880
7881        assert!(
7882            !client.is_active(),
7883            "Client should not be active after idle timeout"
7884        );
7885
7886        client.disconnect().await;
7887        server.abort();
7888    }
7889
7890    #[rstest]
7891    #[tokio::test]
7892    async fn test_idle_timeout_resets_on_data() {
7893        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7894        let port = listener.local_addr().unwrap().port();
7895
7896        // Server sends a message every 200ms (well within 1s idle timeout)
7897        let server = task::spawn(async move {
7898            let (stream, _) = listener.accept().await.unwrap();
7899            let mut ws = accept_async(stream).await.unwrap();
7900
7901            for _ in 0..10 {
7902                sleep(Duration::from_millis(200)).await;
7903
7904                if ws.send(WsMessage::Text("ping".into())).await.is_err() {
7905                    break;
7906                }
7907            }
7908        });
7909
7910        let (handler, _rx) = channel_message_handler();
7911
7912        let config = WebSocketConfig {
7913            url: format!("ws://127.0.0.1:{port}"),
7914            headers: vec![],
7915            heartbeat_interval_secs: None,
7916            heartbeat_payload: None,
7917            connect_timeout_ms: Some(2_000),
7918            reconnect_delay_initial_ms: Some(50),
7919            reconnect_delay_max_ms: Some(100),
7920            reconnect_backoff_factor: Some(1.0),
7921            reconnect_jitter_ms: Some(0),
7922            reconnect_max_attempts: Some(1),
7923            heartbeat_timeout_secs: None,
7924            idle_timeout_ms: Some(1_000),
7925            writer_capacity: None,
7926            backend: TransportBackend::Tungstenite,
7927            proxy_url: None,
7928            max_message_size_bytes: None,
7929            max_frame_size_bytes: None,
7930        };
7931
7932        let client = WebSocketClient::builder()
7933            .config(config)
7934            .message_handler(handler)
7935            .connect()
7936            .await
7937            .unwrap();
7938
7939        assert!(client.is_active());
7940
7941        // Wait 1.5s - data arrives every 200ms so idle timeout (1s) should NOT fire
7942        sleep(Duration::from_millis(1_500)).await;
7943
7944        assert!(
7945            client.is_active(),
7946            "Client should remain active when data is flowing"
7947        );
7948
7949        client.disconnect().await;
7950        server.abort();
7951    }
7952
7953    #[rstest]
7954    #[tokio::test]
7955    async fn test_idle_timeout_fires_when_only_pings_received() {
7956        // Regression: pings and pongs are keep-alive frames, not application data,
7957        // so a peer that only emits control frames must still trip the idle timeout.
7958        // The peer keeps pinging for well past the observation window so the
7959        // pre-fix behavior (reset-on-ping) would keep the client active; under the
7960        // fix the idle timer never resets and fires after ~500ms.
7961        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
7962        let port = listener.local_addr().unwrap().port();
7963        let server = task::spawn(async move {
7964            let (stream, _) = listener.accept().await.unwrap();
7965            let mut ws = accept_async(stream).await.unwrap();
7966
7967            for _ in 0..60 {
7968                sleep(Duration::from_millis(100)).await;
7969
7970                if ws.send(WsMessage::Ping(Vec::new().into())).await.is_err() {
7971                    break;
7972                }
7973            }
7974        });
7975
7976        let (handler, _rx) = channel_message_handler();
7977
7978        let config = WebSocketConfig {
7979            url: format!("ws://127.0.0.1:{port}"),
7980            headers: vec![],
7981            heartbeat_interval_secs: None,
7982            heartbeat_payload: None,
7983            connect_timeout_ms: Some(2_000),
7984            reconnect_delay_initial_ms: Some(50),
7985            reconnect_delay_max_ms: Some(100),
7986            reconnect_backoff_factor: Some(1.0),
7987            reconnect_jitter_ms: Some(0),
7988            reconnect_max_attempts: Some(1),
7989            heartbeat_timeout_secs: None,
7990            idle_timeout_ms: Some(500),
7991            writer_capacity: None,
7992            backend: TransportBackend::Tungstenite,
7993            proxy_url: None,
7994            max_message_size_bytes: None,
7995            max_frame_size_bytes: None,
7996        };
7997
7998        let client = WebSocketClient::builder()
7999            .config(config)
8000            .message_handler(handler)
8001            .connect()
8002            .await
8003            .unwrap();
8004
8005        assert!(client.is_active());
8006
8007        // Observation window is shorter than the ping stream (6s). If the idle
8008        // timer mistakenly reset on every ping the client would still be active
8009        // here; under the fix it goes inactive at ~500ms.
8010        wait_until_async(
8011            || async { client.is_reconnecting() || client.is_disconnected() },
8012            Duration::from_millis(1_500),
8013        )
8014        .await;
8015
8016        assert!(
8017            !client.is_active(),
8018            "Client should not be active after idle timeout when only pings/pongs flow"
8019        );
8020
8021        client.disconnect().await;
8022        server.abort();
8023    }
8024
8025    #[rstest]
8026    #[tokio::test]
8027    async fn test_idle_timeout_fires_when_only_pongs_received() {
8028        // Regression for the heartbeat-reply path. When the client heartbeat is
8029        // enabled, the peer auto-replies with pongs for every outgoing ping. If
8030        // those pongs refreshed last_data_time the idle timer would never fire on
8031        // a zombie connection (the motivating Polymarket scenario).
8032        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
8033        let port = listener.local_addr().unwrap().port();
8034        let server = task::spawn(async move {
8035            let (stream, _) = listener.accept().await.unwrap();
8036            let mut ws = accept_async(stream).await.unwrap();
8037
8038            // Drain incoming frames so tungstenite's internal pong replies are
8039            // actually flushed to the client. Hold the connection open well past
8040            // the observation window.
8041            let deadline = tokio::time::Instant::now() + Duration::from_secs(6);
8042            while tokio::time::Instant::now() < deadline {
8043                if let Ok(Some(Err(_)) | None) =
8044                    tokio::time::timeout(Duration::from_millis(100), ws.next()).await
8045                {
8046                    break;
8047                }
8048            }
8049        });
8050
8051        let (handler, _rx) = channel_message_handler();
8052
8053        let config = WebSocketConfig {
8054            url: format!("ws://127.0.0.1:{port}"),
8055            headers: vec![],
8056            heartbeat_interval_secs: Some(1),
8057            heartbeat_payload: None,
8058            connect_timeout_ms: Some(2_000),
8059            reconnect_delay_initial_ms: Some(50),
8060            reconnect_delay_max_ms: Some(100),
8061            reconnect_backoff_factor: Some(1.0),
8062            reconnect_jitter_ms: Some(0),
8063            reconnect_max_attempts: Some(1),
8064            heartbeat_timeout_secs: None,
8065            idle_timeout_ms: Some(1_500),
8066            writer_capacity: None,
8067            backend: TransportBackend::Tungstenite,
8068            proxy_url: None,
8069            max_message_size_bytes: None,
8070            max_frame_size_bytes: None,
8071        };
8072
8073        let client = WebSocketClient::builder()
8074            .config(config)
8075            .message_handler(handler)
8076            .connect()
8077            .await
8078            .unwrap();
8079
8080        assert!(client.is_active());
8081
8082        // Heartbeat cadence is 1s; each ping draws a pong reply. Under the fix
8083        // the idle timer ignores those pongs and fires at ~1.5s. Under the bug
8084        // every pong reset the timer and the client would stay active.
8085        wait_until_async(
8086            || async { client.is_reconnecting() || client.is_disconnected() },
8087            Duration::from_millis(2_500),
8088        )
8089        .await;
8090
8091        assert!(
8092            !client.is_active(),
8093            "Client should not be active after idle timeout when only pongs flow"
8094        );
8095
8096        client.disconnect().await;
8097        server.abort();
8098    }
8099
8100    #[rstest]
8101    #[tokio::test]
8102    async fn test_disconnect_during_backoff_exits_promptly() {
8103        // Verify that disconnect interrupts backoff sleep (Finding 1).
8104        // Server accepts then drops, no second listener -> reconnect fails -> enters backoff.
8105        // We disconnect while backing off and assert the client shuts down quickly.
8106        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
8107        let port = listener.local_addr().unwrap().port();
8108
8109        let server = task::spawn(async move {
8110            // Accept first connection, close immediately
8111            if let Ok((stream, _)) = listener.accept().await {
8112                let _ = accept_async(stream).await;
8113            }
8114
8115            // Don't accept again so reconnect fails and enters backoff
8116            sleep(Duration::from_mins(1)).await;
8117        });
8118
8119        let (handler, _rx) = channel_message_handler();
8120
8121        let config = WebSocketConfig {
8122            url: format!("ws://127.0.0.1:{port}"),
8123            headers: vec![],
8124            heartbeat_interval_secs: None,
8125            heartbeat_payload: None,
8126            connect_timeout_ms: Some(1_000),
8127            reconnect_delay_initial_ms: Some(10_000), // 10s backoff to ensure we're sleeping
8128            reconnect_delay_max_ms: Some(10_000),
8129            reconnect_backoff_factor: Some(1.0),
8130            reconnect_jitter_ms: Some(0),
8131            reconnect_max_attempts: None,
8132            heartbeat_timeout_secs: None,
8133            idle_timeout_ms: None,
8134            writer_capacity: None,
8135            backend: TransportBackend::Tungstenite,
8136            proxy_url: None,
8137            max_message_size_bytes: None,
8138            max_frame_size_bytes: None,
8139        };
8140
8141        let client = WebSocketClient::builder()
8142            .config(config)
8143            .message_handler(handler)
8144            .connect()
8145            .await
8146            .unwrap();
8147
8148        // Wait for client to enter reconnect
8149        wait_until_async(
8150            || async { client.is_reconnecting() },
8151            Duration::from_secs(3),
8152        )
8153        .await;
8154
8155        // Wait a bit more for the reconnect attempt to fail and enter backoff sleep
8156        sleep(Duration::from_millis(1_500)).await;
8157
8158        // Disconnect while backing off
8159        let start = std::time::Instant::now();
8160        client.disconnect().await;
8161        let elapsed = start.elapsed();
8162
8163        assert!(client.is_disconnected(), "Client should be disconnected");
8164        // Should exit well before the 10s backoff sleep completes
8165        assert!(
8166            elapsed < Duration::from_secs(2),
8167            "Disconnect should interrupt backoff sleep, took {elapsed:?}"
8168        );
8169
8170        server.abort();
8171    }
8172
8173    #[rstest]
8174    #[tokio::test]
8175    async fn test_rate_limit_cancelled_on_disconnect() {
8176        // Verify that a send blocked on rate limiting returns Closed when
8177        // the client disconnects (Finding 6).
8178        use std::{num::NonZeroU32, sync::Arc};
8179
8180        use crate::ratelimiter::quota::Quota;
8181
8182        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
8183        let port = listener.local_addr().unwrap().port();
8184        let server = task::spawn(async move {
8185            if let Ok((stream, _)) = listener.accept().await {
8186                let mut ws = accept_async(stream).await.unwrap();
8187                // Keep alive and echo
8188                while let Some(Ok(msg)) = ws.next().await {
8189                    if ws.send(msg).await.is_err() {
8190                        break;
8191                    }
8192                }
8193            }
8194        });
8195
8196        let (handler, _rx) = channel_message_handler();
8197
8198        let config = WebSocketConfig {
8199            url: format!("ws://127.0.0.1:{port}"),
8200            headers: vec![],
8201            heartbeat_interval_secs: None,
8202            heartbeat_payload: None,
8203            connect_timeout_ms: Some(5_000),
8204            reconnect_delay_initial_ms: Some(100),
8205            reconnect_delay_max_ms: Some(500),
8206            reconnect_backoff_factor: Some(1.5),
8207            reconnect_jitter_ms: Some(0),
8208            reconnect_max_attempts: None,
8209            heartbeat_timeout_secs: None,
8210            idle_timeout_ms: None,
8211            writer_capacity: None,
8212            backend: TransportBackend::Tungstenite,
8213            proxy_url: None,
8214            max_message_size_bytes: None,
8215            max_frame_size_bytes: None,
8216        };
8217
8218        // Very restrictive: 1 req per 60 seconds
8219        let quota = Quota::with_period(Duration::from_mins(1))
8220            .unwrap()
8221            .allow_burst(NonZeroU32::new(1).unwrap());
8222
8223        let client = Arc::new(
8224            WebSocketClient::builder()
8225                .config(config)
8226                .message_handler(handler)
8227                .keyed_quotas(vec![("rate_key".to_string(), quota)])
8228                .connect()
8229                .await
8230                .unwrap(),
8231        );
8232
8233        let test_key: [Ustr; 1] = [Ustr::from("rate_key")];
8234
8235        // Exhaust the burst quota
8236        client
8237            .send_text("exhaust".to_string(), Some(test_key.as_slice()))
8238            .await
8239            .unwrap();
8240
8241        // Spawn a send that will block on rate limiter
8242        let client_clone = Arc::clone(&client);
8243        let send_handle = task::spawn(async move {
8244            client_clone
8245                .send_text("blocked".to_string(), Some(&[Ustr::from("rate_key")]))
8246                .await
8247        });
8248
8249        // Let the send block on rate limit
8250        sleep(Duration::from_millis(200)).await;
8251
8252        // Disconnect while send is blocked
8253        let start = std::time::Instant::now();
8254        client.disconnect().await;
8255        let elapsed_disconnect = start.elapsed();
8256
8257        // The blocked send should return Closed
8258        let result = tokio::time::timeout(Duration::from_secs(2), send_handle)
8259            .await
8260            .expect("Send task should complete quickly")
8261            .expect("Send task should not panic");
8262
8263        assert!(
8264            matches!(result, Err(crate::error::SendError::Closed)),
8265            "Blocked send should return Closed, was: {result:?}"
8266        );
8267
8268        // Disconnect should be fast, not waiting for the 60s rate limit
8269        assert!(
8270            elapsed_disconnect < Duration::from_secs(3),
8271            "Disconnect should not wait for rate limiter, took {elapsed_disconnect:?}"
8272        );
8273
8274        server.abort();
8275    }
8276
8277    #[rstest]
8278    #[tokio::test]
8279    async fn test_stream_mode_transitions_to_closed_on_dead_write_task() {
8280        // Verify that stream mode transitions to CLOSED (not RECONNECT) when
8281        // the write task dies (Finding 4). We force write failure by sending
8282        // after the server closes the connection.
8283        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
8284        let port = listener.local_addr().unwrap().port();
8285        let server = task::spawn(async move {
8286            if let Ok((stream, _)) = listener.accept().await
8287                && let Ok(ws) = accept_async(stream).await
8288            {
8289                // Close immediately to cause write errors
8290                drop(ws);
8291            }
8292        });
8293
8294        let config = WebSocketConfig {
8295            url: format!("ws://127.0.0.1:{port}"),
8296            headers: vec![],
8297            heartbeat_interval_secs: None,
8298            heartbeat_payload: None,
8299            connect_timeout_ms: Some(1_000),
8300            reconnect_delay_initial_ms: Some(50),
8301            reconnect_delay_max_ms: Some(100),
8302            reconnect_backoff_factor: Some(1.0),
8303            reconnect_jitter_ms: Some(0),
8304            reconnect_max_attempts: None,
8305            heartbeat_timeout_secs: None,
8306            idle_timeout_ms: None,
8307            writer_capacity: None,
8308            backend: TransportBackend::Tungstenite,
8309            proxy_url: None,
8310            max_message_size_bytes: None,
8311            max_frame_size_bytes: None,
8312        };
8313
8314        let (_reader, client) = WebSocketClient::stream_builder()
8315            .config(config)
8316            .connect()
8317            .await
8318            .unwrap();
8319
8320        assert!(client.is_active(), "Client should start active");
8321
8322        // Wait for server to close, then send to trigger write task failure
8323        sleep(Duration::from_millis(100)).await;
8324
8325        // Keep sending until the write task detects the broken connection
8326        for _ in 0..20 {
8327            let _ = client.send_text("ping".to_string(), None).await;
8328            sleep(Duration::from_millis(50)).await;
8329
8330            if !client.is_active() {
8331                break;
8332            }
8333        }
8334
8335        // Wait for controller to process the state change
8336        wait_until_async(|| async { !client.is_active() }, Duration::from_secs(5)).await;
8337
8338        // Stream mode should go to CLOSED, not RECONNECT
8339        assert!(
8340            client.is_closed() || client.is_disconnected(),
8341            "Stream mode should transition to CLOSED, not RECONNECT. \
8342             is_reconnecting={}, is_closed={}, is_disconnected={}",
8343            client.is_reconnecting(),
8344            client.is_closed(),
8345            client.is_disconnected(),
8346        );
8347        assert!(
8348            !client.is_reconnecting(),
8349            "Stream mode should never attempt reconnection"
8350        );
8351
8352        server.abort();
8353    }
8354
8355    #[derive(Default)]
8356    struct BlockingFailState {
8357        send_entered: AtomicBool,
8358        send_entered_notify: tokio::sync::Notify,
8359        released: AtomicBool,
8360        fail: AtomicBool,
8361        waker: parking_lot::Mutex<Option<std::task::Waker>>,
8362    }
8363
8364    impl BlockingFailState {
8365        fn trigger_failure(&self) {
8366            self.fail.store(true, Ordering::SeqCst);
8367            self.release_send();
8368        }
8369
8370        fn release_send(&self) {
8371            self.released.store(true, Ordering::SeqCst);
8372
8373            if let Some(waker) = self.waker.lock().take() {
8374                waker.wake();
8375            }
8376        }
8377    }
8378
8379    /// Transport whose sends block until released and can fail on demand.
8380    struct BlockingFailTransport {
8381        state: Arc<BlockingFailState>,
8382    }
8383
8384    impl futures_util::Stream for BlockingFailTransport {
8385        type Item = Result<Message, TransportError>;
8386
8387        fn poll_next(
8388            self: std::pin::Pin<&mut Self>,
8389            _cx: &mut std::task::Context<'_>,
8390        ) -> std::task::Poll<Option<Self::Item>> {
8391            std::task::Poll::Pending
8392        }
8393    }
8394
8395    impl futures_util::Sink<Message> for BlockingFailTransport {
8396        type Error = TransportError;
8397
8398        fn poll_ready(
8399            self: std::pin::Pin<&mut Self>,
8400            _cx: &mut std::task::Context<'_>,
8401        ) -> std::task::Poll<Result<(), Self::Error>> {
8402            std::task::Poll::Ready(Ok(()))
8403        }
8404
8405        fn start_send(self: std::pin::Pin<&mut Self>, _item: Message) -> Result<(), Self::Error> {
8406            Ok(())
8407        }
8408
8409        fn poll_flush(
8410            self: std::pin::Pin<&mut Self>,
8411            cx: &mut std::task::Context<'_>,
8412        ) -> std::task::Poll<Result<(), Self::Error>> {
8413            // Store the waker before checking the flag so trigger_failure
8414            // cannot slip between the check and the registration
8415            *self.state.waker.lock() = Some(cx.waker().clone());
8416            self.state.send_entered.store(true, Ordering::SeqCst);
8417            self.state.send_entered_notify.notify_one();
8418
8419            if !self.state.released.load(Ordering::SeqCst) {
8420                std::task::Poll::Pending
8421            } else if self.state.fail.load(Ordering::SeqCst) {
8422                std::task::Poll::Ready(Err(TransportError::ConnectionReset))
8423            } else {
8424                std::task::Poll::Ready(Ok(()))
8425            }
8426        }
8427
8428        fn poll_close(
8429            self: std::pin::Pin<&mut Self>,
8430            _cx: &mut std::task::Context<'_>,
8431        ) -> std::task::Poll<Result<(), Self::Error>> {
8432            std::task::Poll::Ready(Ok(()))
8433        }
8434    }
8435
8436    struct BlockingMessageState {
8437        polled_tx: Mutex<Option<std::sync::mpsc::Sender<()>>>,
8438        release: (Mutex<bool>, parking_lot::Condvar),
8439        message: Mutex<Option<Message>>,
8440    }
8441
8442    struct BlockingMessageTransport {
8443        state: Arc<BlockingMessageState>,
8444    }
8445
8446    impl futures_util::Stream for BlockingMessageTransport {
8447        type Item = Result<Message, TransportError>;
8448
8449        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
8450            if let Some(polled_tx) = self.state.polled_tx.lock().take() {
8451                polled_tx.send(()).unwrap();
8452            }
8453
8454            let (lock, condvar) = &self.state.release;
8455            let mut released = lock.lock();
8456            while !*released {
8457                condvar.wait(&mut released);
8458            }
8459
8460            Poll::Ready(self.state.message.lock().take().map(Ok))
8461        }
8462    }
8463
8464    impl futures_util::Sink<Message> for BlockingMessageTransport {
8465        type Error = TransportError;
8466
8467        fn poll_ready(
8468            self: Pin<&mut Self>,
8469            _cx: &mut Context<'_>,
8470        ) -> Poll<Result<(), Self::Error>> {
8471            Poll::Ready(Ok(()))
8472        }
8473
8474        fn start_send(self: Pin<&mut Self>, _item: Message) -> Result<(), Self::Error> {
8475            Ok(())
8476        }
8477
8478        fn poll_flush(
8479            self: Pin<&mut Self>,
8480            _cx: &mut Context<'_>,
8481        ) -> Poll<Result<(), Self::Error>> {
8482            Poll::Ready(Ok(()))
8483        }
8484
8485        fn poll_close(
8486            self: Pin<&mut Self>,
8487            _cx: &mut Context<'_>,
8488        ) -> Poll<Result<(), Self::Error>> {
8489            Poll::Ready(Ok(()))
8490        }
8491    }
8492
8493    struct RecordingState {
8494        messages: Arc<Mutex<Vec<Message>>>,
8495        recorded_notify: tokio::sync::Notify,
8496    }
8497
8498    struct RecordingTransport {
8499        state: Arc<RecordingState>,
8500    }
8501
8502    impl futures_util::Stream for RecordingTransport {
8503        type Item = Result<Message, TransportError>;
8504
8505        fn poll_next(self: Pin<&mut Self>, _cx: &mut Context<'_>) -> Poll<Option<Self::Item>> {
8506            Poll::Pending
8507        }
8508    }
8509
8510    impl futures_util::Sink<Message> for RecordingTransport {
8511        type Error = TransportError;
8512
8513        fn poll_ready(
8514            self: Pin<&mut Self>,
8515            _cx: &mut Context<'_>,
8516        ) -> Poll<Result<(), Self::Error>> {
8517            Poll::Ready(Ok(()))
8518        }
8519
8520        fn start_send(self: Pin<&mut Self>, item: Message) -> Result<(), Self::Error> {
8521            self.state.messages.lock().push(item);
8522            self.state.recorded_notify.notify_one();
8523            Ok(())
8524        }
8525
8526        fn poll_flush(
8527            self: Pin<&mut Self>,
8528            _cx: &mut Context<'_>,
8529        ) -> Poll<Result<(), Self::Error>> {
8530            Poll::Ready(Ok(()))
8531        }
8532
8533        fn poll_close(
8534            self: Pin<&mut Self>,
8535            _cx: &mut Context<'_>,
8536        ) -> Poll<Result<(), Self::Error>> {
8537            Poll::Ready(Ok(()))
8538        }
8539    }
8540
8541    #[rstest]
8542    #[tokio::test(start_paused = true)]
8543    async fn test_pong_is_bound_to_connection_epoch() {
8544        let initial_state = Arc::new(RecordingState {
8545            messages: Arc::new(Mutex::new(Vec::new())),
8546            recorded_notify: tokio::sync::Notify::new(),
8547        });
8548
8549        let initial_transport: BoxedWsTransport = Box::pin(RecordingTransport {
8550            state: initial_state,
8551        });
8552
8553        let (writer, _reader) = initial_transport.split();
8554        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8555        let state_notify = Arc::new(tokio::sync::Notify::new());
8556        let auth_tracker = Arc::new(OnceLock::new());
8557        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
8558        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
8559        let write_task = WebSocketClientInner::spawn_write_task(
8560            Arc::clone(&connection_state),
8561            Arc::clone(&state_notify),
8562            Arc::new(AtomicBool::new(true)),
8563            writer,
8564            writer_rx,
8565            Arc::new(AtomicU64::new(0)),
8566            auth_tracker,
8567            reconnect_buffer_waits_for_auth,
8568            None,
8569        );
8570
8571        let recorded = Arc::new(Mutex::new(Vec::new()));
8572
8573        let replacement_state = Arc::new(RecordingState {
8574            messages: Arc::clone(&recorded),
8575            recorded_notify: tokio::sync::Notify::new(),
8576        });
8577
8578        let replacement_transport: BoxedWsTransport = Box::pin(RecordingTransport {
8579            state: replacement_state,
8580        });
8581
8582        let (replacement_writer, _reader) = replacement_transport.split();
8583        let (update_tx, update_rx) = tokio::sync::oneshot::channel();
8584        writer_tx
8585            .send(WriterCommand::Update(replacement_writer, update_tx))
8586            .unwrap();
8587        writer_tx
8588            .send_control(WriterCommand::SendPongOnConnection {
8589                data: b"stale-pong".to_vec(),
8590                connection_epoch: 0,
8591            })
8592            .unwrap();
8593
8594        tokio::time::advance(Duration::from_millis(100)).await;
8595        assert_eq!(update_rx.await.unwrap(), 1);
8596
8597        let (sentinel_tx, sentinel_rx) = tokio::sync::oneshot::channel();
8598        writer_tx
8599            .send_control(WriterCommand::SendOnConnection {
8600                message: Message::text("sentinel-1"),
8601                connection_epoch: 1,
8602                response_tx: sentinel_tx,
8603            })
8604            .unwrap();
8605
8606        sentinel_rx.await.unwrap().unwrap();
8607        assert_eq!(recorded.lock().as_slice(), &[Message::text("sentinel-1")]);
8608
8609        writer_tx
8610            .send_control(WriterCommand::SendPongOnConnection {
8611                data: b"fresh-pong".to_vec(),
8612                connection_epoch: 1,
8613            })
8614            .unwrap();
8615
8616        let (sentinel_tx, sentinel_rx) = tokio::sync::oneshot::channel();
8617        writer_tx
8618            .send_control(WriterCommand::SendOnConnection {
8619                message: Message::text("sentinel-2"),
8620                connection_epoch: 1,
8621                response_tx: sentinel_tx,
8622            })
8623            .unwrap();
8624
8625        sentinel_rx.await.unwrap().unwrap();
8626        assert_eq!(
8627            recorded.lock().as_slice(),
8628            &[
8629                Message::text("sentinel-1"),
8630                Message::Pong(b"fresh-pong".to_vec().into()),
8631                Message::text("sentinel-2"),
8632            ]
8633        );
8634
8635        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
8636        state_notify.notify_waiters();
8637        drop(writer_tx);
8638        write_task.await.unwrap();
8639    }
8640
8641    #[rstest]
8642    #[case(Message::text("stale"))]
8643    #[case(Message::Binary(vec![1, 2, 3].into()))]
8644    #[case(Message::Ping(vec![1, 2, 3].into()))]
8645    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
8646    async fn test_message_handler_drops_old_session_message(#[case] message: Message) {
8647        let (polled_tx, polled_rx) = std::sync::mpsc::channel();
8648
8649        let state = Arc::new(BlockingMessageState {
8650            polled_tx: Mutex::new(Some(polled_tx)),
8651            release: (Mutex::new(false), Condvar::new()),
8652            message: Mutex::new(Some(message)),
8653        });
8654
8655        let release_guard = CondvarReleaseGuard::new(&state.release);
8656
8657        let transport: BoxedWsTransport = Box::pin(BlockingMessageTransport {
8658            state: Arc::clone(&state),
8659        });
8660
8661        let (_writer, reader) = transport.split();
8662        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8663        let state_notify = Arc::new(tokio::sync::Notify::new());
8664        let read_fence = ReadSessionFence::new();
8665        let message_count = Arc::new(AtomicUsize::new(0));
8666        let ping_count = Arc::new(AtomicUsize::new(0));
8667        let message_count_clone = Arc::clone(&message_count);
8668        let ping_count_clone = Arc::clone(&ping_count);
8669        let message_handler: MessageHandler =
8670            Arc::new(move |_| _ = message_count_clone.fetch_add(1, Ordering::SeqCst));
8671        let message_handler = IncomingHandler::Message(message_handler);
8672        let ping_handler: PingHandler =
8673            Arc::new(move |_| _ = ping_count_clone.fetch_add(1, Ordering::SeqCst));
8674        let ping_handler = IncomingPingHandler::Ping(ping_handler);
8675
8676        let read_task = WebSocketClientInner::spawn_message_handler_task(
8677            Arc::clone(&connection_state),
8678            state_notify,
8679            read_fence.clone(),
8680            reader,
8681            0,
8682            Some(&message_handler),
8683            Some(&ping_handler),
8684            None,
8685            None,
8686            None,
8687            None,
8688        );
8689
8690        recv_rendezvous(polled_rx, "WebSocket reader poll entry").await;
8691        connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
8692        read_fence.invalidate();
8693        read_task.abort();
8694        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
8695        release_guard.release();
8696        await_task_termination(read_task, "old WebSocket read task").await;
8697
8698        assert_eq!(message_count.load(Ordering::SeqCst), 0);
8699        assert_eq!(ping_count.load(Ordering::SeqCst), 0);
8700    }
8701
8702    #[rstest]
8703    #[tokio::test]
8704    async fn test_reconnect_buffer_drain_stops_after_reconnect_request() {
8705        let state = Arc::new(BlockingFailState::default());
8706
8707        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
8708            state: Arc::clone(&state),
8709        });
8710
8711        let (mut writer, _reader) = transport.split();
8712        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8713        let task_connection_state = Arc::clone(&connection_state);
8714
8715        let drain_task = tokio::spawn(async move {
8716            let auth_tracker = Arc::new(OnceLock::new());
8717            let reconnect_buffer_waits_for_auth = AtomicBool::new(false);
8718            let (sender, mut receiver) = crate::writer::channel(2);
8719            sender.send(Message::text("admitted")).unwrap();
8720            sender.send(Message::text("held-for-reconnect")).unwrap();
8721            let mut buffer = VecDeque::from([
8722                receiver.recv().await.unwrap(),
8723                receiver.recv().await.unwrap(),
8724            ]);
8725            let send_error = WebSocketClientInner::drain_reconnect_buffer(
8726                &mut buffer,
8727                &mut writer,
8728                &task_connection_state,
8729                &auth_tracker,
8730                &reconnect_buffer_waits_for_auth,
8731            )
8732            .await;
8733            (buffer, send_error)
8734        });
8735
8736        state.send_entered_notify.notified().await;
8737        connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
8738        state.release_send();
8739
8740        let (buffer, send_error) = tokio::time::timeout(TEST_TIMEOUT, drain_task)
8741            .await
8742            .expect("buffer drain should stop after reconnect acceptance")
8743            .unwrap();
8744        assert!(!send_error);
8745        assert_eq!(
8746            buffer
8747                .into_iter()
8748                .map(|(message, _)| message)
8749                .collect::<Vec<_>>(),
8750            vec![Message::text("held-for-reconnect")]
8751        );
8752    }
8753
8754    #[rstest]
8755    #[case(0)]
8756    #[case(usize::MAX)]
8757    #[tokio::test]
8758    async fn test_stream_writer_rejects_invalid_capacity(#[case] capacity: usize) {
8759        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
8760            state: Arc::new(BlockingFailState::default()),
8761        });
8762
8763        let (writer, _reader) = transport.split();
8764        let mut config = WebSocketConfig::builder()
8765            .url("ws://localhost".to_string())
8766            .build()
8767            .unwrap();
8768        config.writer_capacity = Some(capacity);
8769
8770        let result = WebSocketClientInner::new_with_writer(config, writer).await;
8771
8772        assert!(
8773            matches!(result, Err(TransportError::Io(e)) if e.kind() == std::io::ErrorKind::InvalidInput && e.to_string().contains("writer_capacity"))
8774        );
8775    }
8776
8777    #[rstest]
8778    #[case(0)]
8779    #[case(usize::MAX)]
8780    #[tokio::test]
8781    async fn test_handlerless_connect_rejects_invalid_capacity(#[case] capacity: usize) {
8782        let mut config = WebSocketConfig::builder()
8783            .url("ws://127.0.0.1:1".to_string())
8784            .build()
8785            .unwrap();
8786        config.writer_capacity = Some(capacity);
8787
8788        let result = WebSocketClientInner::connect_url(config, None, None).await;
8789
8790        assert!(
8791            matches!(result, Err(TransportError::Io(e)) if e.kind() == std::io::ErrorKind::InvalidInput && e.to_string().contains("writer_capacity"))
8792        );
8793    }
8794
8795    #[rstest]
8796    #[tokio::test(start_paused = true)]
8797    async fn test_stalled_websocket_send_reconnects_and_replays() {
8798        let state = Arc::new(BlockingFailState::default());
8799
8800        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
8801            state: Arc::clone(&state),
8802        });
8803
8804        let (writer, _reader) = transport.split();
8805        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8806        let state_notify = Arc::new(tokio::sync::Notify::new());
8807        let auth_tracker = Arc::new(OnceLock::new());
8808        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
8809        let (writer_tx, writer_rx) = crate::writer::channel(2);
8810        let states = Arc::new(Mutex::new(Vec::new()));
8811        let states_callback = Arc::clone(&states);
8812        let sink = SocketStateSink::new(move |state| {
8813            states_callback.lock().push(state);
8814        });
8815
8816        let write_task = WebSocketClientInner::spawn_write_task(
8817            Arc::clone(&connection_state),
8818            Arc::clone(&state_notify),
8819            Arc::new(AtomicBool::new(true)),
8820            writer,
8821            writer_rx,
8822            Arc::new(AtomicU64::new(0)),
8823            Arc::clone(&auth_tracker),
8824            reconnect_buffer_waits_for_auth,
8825            Some(sink),
8826        );
8827
8828        writer_tx
8829            .send(WriterCommand::Send(Message::text("complete-message")))
8830            .unwrap();
8831        state.send_entered_notify.notified().await;
8832        writer_tx
8833            .send(WriterCommand::Send(Message::text("second-message")))
8834            .unwrap();
8835        assert_eq!(
8836            ConnectionMode::from_atomic(&connection_state),
8837            ConnectionMode::Active
8838        );
8839
8840        for _ in 0..100 {
8841            assert!(matches!(
8842                writer_tx.send(WriterCommand::Send(Message::text("overflow"))),
8843                Err(SendError::BufferFull)
8844            ));
8845        }
8846
8847        tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
8848        tokio::task::yield_now().await;
8849        assert_eq!(
8850            ConnectionMode::from_atomic(&connection_state),
8851            ConnectionMode::Reconnect
8852        );
8853        assert!(matches!(
8854            writer_tx.send(WriterCommand::Send(Message::text("overflow"))),
8855            Err(SendError::BufferFull)
8856        ));
8857
8858        let recorded = Arc::new(Mutex::new(Vec::new()));
8859
8860        let recording_state = Arc::new(RecordingState {
8861            messages: Arc::clone(&recorded),
8862            recorded_notify: tokio::sync::Notify::new(),
8863        });
8864
8865        let transport: BoxedWsTransport = Box::pin(RecordingTransport {
8866            state: Arc::clone(&recording_state),
8867        });
8868
8869        let (new_writer, _reader) = transport.split();
8870        let (update_tx, update_rx) = tokio::sync::oneshot::channel();
8871        writer_tx
8872            .send_update(WriterCommand::Update(new_writer, update_tx))
8873            .unwrap();
8874
8875        tokio::time::advance(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS)).await;
8876        assert_eq!(
8877            tokio::time::timeout(Duration::from_secs(1), update_rx)
8878                .await
8879                .expect("writer update should not remain queued behind a stalled send")
8880                .unwrap(),
8881            1,
8882            "the replacement sink should install as connection epoch 1"
8883        );
8884        assert_eq!(
8885            ConnectionMode::from_atomic(&connection_state),
8886            ConnectionMode::Reconnect
8887        );
8888
8889        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
8890        state_notify.notify_waiters();
8891        tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
8892        recording_state.recorded_notify.notified().await;
8893        assert_eq!(
8894            recorded.lock().as_slice(),
8895            &[
8896                Message::text("complete-message"),
8897                Message::text("second-message")
8898            ]
8899        );
8900
8901        writer_tx
8902            .send(WriterCommand::Send(Message::text("second-message")))
8903            .unwrap();
8904
8905        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
8906        state_notify.notify_waiters();
8907        drop(writer_tx);
8908        write_task.await.unwrap();
8909
8910        assert_eq!(*states.lock(), vec![SocketState::Disconnected]);
8911    }
8912
8913    #[rstest]
8914    #[case(Message::Ping(vec![1, 2, 3].into()))]
8915    #[case(Message::Pong(vec![4, 5, 6].into()))]
8916    #[case(Message::Close(None))]
8917    #[tokio::test(start_paused = true)]
8918    async fn test_stalled_control_frame_is_not_replayed(#[case] control: Message) {
8919        // A control frame belongs to the connection it was issued on, so a failed write must
8920        // drop it instead of replaying it onto the replacement connection.
8921        let state = Arc::new(BlockingFailState::default());
8922
8923        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
8924            state: Arc::clone(&state),
8925        });
8926
8927        let (writer, _reader) = transport.split();
8928        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
8929        let state_notify = Arc::new(tokio::sync::Notify::new());
8930        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
8931        let write_task = WebSocketClientInner::spawn_write_task(
8932            Arc::clone(&connection_state),
8933            Arc::clone(&state_notify),
8934            Arc::new(AtomicBool::new(true)),
8935            writer,
8936            writer_rx,
8937            Arc::new(AtomicU64::new(0)),
8938            Arc::new(OnceLock::new()),
8939            Arc::new(AtomicBool::new(false)),
8940            None,
8941        );
8942
8943        writer_tx.send(WriterCommand::Send(control)).unwrap();
8944        state.send_entered_notify.notified().await;
8945
8946        let recorded = Arc::new(Mutex::new(Vec::new()));
8947
8948        let recording_state = Arc::new(RecordingState {
8949            messages: Arc::clone(&recorded),
8950            recorded_notify: tokio::sync::Notify::new(),
8951        });
8952
8953        let transport: BoxedWsTransport = Box::pin(RecordingTransport {
8954            state: recording_state,
8955        });
8956
8957        let (new_writer, _reader) = transport.split();
8958        let (update_tx, update_rx) = tokio::sync::oneshot::channel();
8959        writer_tx
8960            .send(WriterCommand::Update(new_writer, update_tx))
8961            .unwrap();
8962
8963        tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
8964        assert_eq!(
8965            tokio::time::timeout(Duration::from_secs(1), update_rx)
8966                .await
8967                .expect("writer update should not remain queued behind a stalled send")
8968                .unwrap(),
8969            1,
8970            "the replacement sink should install as connection epoch 1"
8971        );
8972        assert_eq!(
8973            ConnectionMode::from_atomic(&connection_state),
8974            ConnectionMode::Reconnect,
8975            "a failed control-frame write should still trigger a reconnect"
8976        );
8977
8978        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
8979        state_notify.notify_waiters();
8980
8981        // Each pass returns the writer to the loop top, where a buffered frame would drain
8982        for _ in 0..5 {
8983            tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
8984        }
8985
8986        let replayed = recorded.lock().clone();
8987        assert!(
8988            replayed.is_empty(),
8989            "a failed control frame must not reach the replacement connection, was {replayed:?}"
8990        );
8991
8992        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
8993        state_notify.notify_waiters();
8994        drop(writer_tx);
8995        write_task.await.unwrap();
8996    }
8997
8998    #[rstest]
8999    #[case(Message::Ping(vec![1, 2, 3].into()))]
9000    #[case(Message::Pong(vec![4, 5, 6].into()))]
9001    #[case(Message::Close(None))]
9002    #[tokio::test(start_paused = true)]
9003    async fn test_control_frame_enqueued_during_reconnect_is_not_replayed(
9004        #[case] control: Message,
9005    ) {
9006        // `send_pong` and the heartbeat task check for an active connection before enqueueing,
9007        // so a mode flip can land their frame on the writer's reconnect-mode branch instead.
9008        let recorded = Arc::new(Mutex::new(Vec::new()));
9009
9010        let recording_state = Arc::new(RecordingState {
9011            messages: Arc::clone(&recorded),
9012            recorded_notify: tokio::sync::Notify::new(),
9013        });
9014
9015        let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9016            state: recording_state,
9017        });
9018
9019        let (writer, _reader) = transport.split();
9020        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
9021        let state_notify = Arc::new(tokio::sync::Notify::new());
9022        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9023        let write_task = WebSocketClientInner::spawn_write_task(
9024            Arc::clone(&connection_state),
9025            Arc::clone(&state_notify),
9026            Arc::new(AtomicBool::new(true)),
9027            writer,
9028            writer_rx,
9029            Arc::new(AtomicU64::new(0)),
9030            Arc::new(OnceLock::new()),
9031            Arc::new(AtomicBool::new(false)),
9032            None,
9033        );
9034
9035        writer_tx.send(WriterCommand::Send(control)).unwrap();
9036        tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9037
9038        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9039        state_notify.notify_waiters();
9040
9041        // Each pass returns the writer to the loop top, where a buffered frame would drain
9042        for _ in 0..5 {
9043            tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9044        }
9045
9046        let replayed = recorded.lock().clone();
9047        assert!(
9048            replayed.is_empty(),
9049            "a control frame enqueued during reconnect must not reach the replacement connection, was {replayed:?}"
9050        );
9051
9052        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9053        state_notify.notify_waiters();
9054        drop(writer_tx);
9055        write_task.await.unwrap();
9056    }
9057
9058    #[rstest]
9059    #[tokio::test(start_paused = true)]
9060    async fn test_stalled_text_heartbeat_is_not_replayed() {
9061        let state = Arc::new(BlockingFailState::default());
9062
9063        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9064            state: Arc::clone(&state),
9065        });
9066
9067        let (writer, _reader) = transport.split();
9068        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9069        let state_notify = Arc::new(tokio::sync::Notify::new());
9070        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9071        let write_task = WebSocketClientInner::spawn_write_task(
9072            Arc::clone(&connection_state),
9073            Arc::clone(&state_notify),
9074            Arc::new(AtomicBool::new(true)),
9075            writer,
9076            writer_rx,
9077            Arc::new(AtomicU64::new(0)),
9078            Arc::new(OnceLock::new()),
9079            Arc::new(AtomicBool::new(false)),
9080            None,
9081        );
9082
9083        writer_tx
9084            .send(WriterCommand::Heartbeat(Message::text(
9085                "{\"op\":\"heartbeat\"}",
9086            )))
9087            .unwrap();
9088        state.send_entered_notify.notified().await;
9089
9090        let recorded = Arc::new(Mutex::new(Vec::new()));
9091
9092        let recording_state = Arc::new(RecordingState {
9093            messages: Arc::clone(&recorded),
9094            recorded_notify: tokio::sync::Notify::new(),
9095        });
9096
9097        let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9098            state: recording_state,
9099        });
9100
9101        let (new_writer, _reader) = transport.split();
9102        let (update_tx, update_rx) = tokio::sync::oneshot::channel();
9103        writer_tx
9104            .send(WriterCommand::Update(new_writer, update_tx))
9105            .unwrap();
9106
9107        tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
9108        assert_eq!(
9109            tokio::time::timeout(Duration::from_secs(1), update_rx)
9110                .await
9111                .expect("writer update should not remain queued behind a stalled send")
9112                .unwrap(),
9113            1,
9114            "the replacement sink should install as connection epoch 1"
9115        );
9116        assert_eq!(
9117            ConnectionMode::from_atomic(&connection_state),
9118            ConnectionMode::Reconnect,
9119            "a failed heartbeat write should still trigger a reconnect"
9120        );
9121
9122        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9123        state_notify.notify_waiters();
9124
9125        for _ in 0..5 {
9126            tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9127        }
9128
9129        let replayed = recorded.lock().clone();
9130        assert!(
9131            replayed.is_empty(),
9132            "a failed text heartbeat must not reach the replacement connection, was {replayed:?}"
9133        );
9134
9135        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9136        state_notify.notify_waiters();
9137        drop(writer_tx);
9138        write_task.await.unwrap();
9139    }
9140
9141    #[rstest]
9142    #[tokio::test(start_paused = true)]
9143    async fn test_text_heartbeat_enqueued_during_reconnect_is_not_replayed() {
9144        let recorded = Arc::new(Mutex::new(Vec::new()));
9145
9146        let recording_state = Arc::new(RecordingState {
9147            messages: Arc::clone(&recorded),
9148            recorded_notify: tokio::sync::Notify::new(),
9149        });
9150
9151        let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9152            state: recording_state,
9153        });
9154
9155        let (writer, _reader) = transport.split();
9156        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
9157        let state_notify = Arc::new(tokio::sync::Notify::new());
9158        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9159        let write_task = WebSocketClientInner::spawn_write_task(
9160            Arc::clone(&connection_state),
9161            Arc::clone(&state_notify),
9162            Arc::new(AtomicBool::new(true)),
9163            writer,
9164            writer_rx,
9165            Arc::new(AtomicU64::new(0)),
9166            Arc::new(OnceLock::new()),
9167            Arc::new(AtomicBool::new(false)),
9168            None,
9169        );
9170
9171        writer_tx
9172            .send(WriterCommand::Heartbeat(Message::text(
9173                "{\"op\":\"heartbeat\"}",
9174            )))
9175            .unwrap();
9176        tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9177
9178        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9179        state_notify.notify_waiters();
9180
9181        for _ in 0..5 {
9182            tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9183        }
9184
9185        let replayed = recorded.lock().clone();
9186        assert!(
9187            replayed.is_empty(),
9188            "a text heartbeat enqueued during reconnect must not reach the replacement connection, was {replayed:?}"
9189        );
9190
9191        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9192        state_notify.notify_waiters();
9193        drop(writer_tx);
9194        write_task.await.unwrap();
9195    }
9196
9197    #[rstest]
9198    #[case::text(
9199        Some("{\"op\":\"heartbeat\"}"),
9200        Message::text("{\"op\":\"heartbeat\"}")
9201    )]
9202    #[case::ping(None, Message::Ping(vec![].into()))]
9203    #[tokio::test(start_paused = true)]
9204    async fn test_heartbeat_task_enqueues_writer_heartbeat_command(
9205        #[case] payload: Option<&str>,
9206        #[case] expected: Message,
9207    ) {
9208        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9209        let (writer_tx, mut writer_rx) =
9210            crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9211        let task = WebSocketClientInner::spawn_heartbeat_task(
9212            Arc::clone(&connection_state),
9213            1,
9214            payload.map(ToString::to_string),
9215            writer_tx,
9216        );
9217
9218        tokio::time::advance(Duration::from_secs(1)).await;
9219        let cmd = writer_rx
9220            .recv()
9221            .await
9222            .expect("heartbeat task should enqueue");
9223
9224        match cmd.0 {
9225            WriterCommand::Heartbeat(msg) => assert_eq!(msg, expected),
9226            other => panic!("expected Heartbeat, was {other:?}"),
9227        }
9228
9229        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9230        tokio::time::advance(Duration::from_secs(1)).await;
9231        task.await.unwrap();
9232    }
9233
9234    #[cfg(target_os = "linux")]
9235    #[rstest]
9236    #[tokio::test]
9237    async fn test_heartbeat_task_warns_when_writer_channel_closed() {
9238        let capture = capture_logs_for(&["nautilus_network::websocket::client"]).await;
9239        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9240        let (writer_tx, mut writer_rx) =
9241            crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9242        let task = WebSocketClientInner::spawn_heartbeat_task(
9243            Arc::clone(&connection_state),
9244            1,
9245            None,
9246            writer_tx,
9247        );
9248
9249        let _ = tokio::time::timeout(Duration::from_secs(2), writer_rx.recv())
9250            .await
9251            .expect("timed out waiting for the first heartbeat")
9252            .expect("heartbeat channel closed before the first heartbeat");
9253        drop(writer_rx);
9254        tokio::time::sleep(Duration::from_millis(1_200)).await;
9255
9256        let messages = capture.messages();
9257        assert!(
9258            messages.iter().any(|(level, message)| {
9259                *level == Level::Warn && message.contains("Failed to send heartbeat to writer task")
9260            }),
9261            "expected WARN when the writer channel is closed, was {messages:?}"
9262        );
9263        assert!(
9264            !messages.iter().any(|(level, _)| *level == Level::Error),
9265            "a closed writer channel must not log ERROR, was {messages:?}"
9266        );
9267
9268        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9269        tokio::time::timeout(Duration::from_secs(2), task)
9270            .await
9271            .expect("heartbeat task should stop after close")
9272            .unwrap();
9273    }
9274
9275    #[rstest]
9276    #[tokio::test(start_paused = true)]
9277    async fn test_stalled_ownership_bound_send_times_out_without_replay() {
9278        let state = Arc::new(BlockingFailState::default());
9279
9280        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9281            state: Arc::clone(&state),
9282        });
9283
9284        let (writer, _reader) = transport.split();
9285        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9286        let state_notify = Arc::new(tokio::sync::Notify::new());
9287        let auth_tracker = Arc::new(OnceLock::new());
9288        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
9289        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9290        let write_task = WebSocketClientInner::spawn_write_task(
9291            Arc::clone(&connection_state),
9292            Arc::clone(&state_notify),
9293            Arc::new(AtomicBool::new(true)),
9294            writer,
9295            writer_rx,
9296            Arc::new(AtomicU64::new(0)),
9297            Arc::clone(&auth_tracker),
9298            reconnect_buffer_waits_for_auth,
9299            None,
9300        );
9301
9302        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9303        writer_tx
9304            .send_control(WriterCommand::SendOnConnection {
9305                message: Message::text("ownership-bound"),
9306                connection_epoch: 0,
9307                response_tx,
9308            })
9309            .unwrap();
9310
9311        state.send_entered_notify.notified().await;
9312
9313        tokio::time::advance(Duration::from_secs(WRITE_TIMEOUT_SECS)).await;
9314        let outcome = tokio::time::timeout(Duration::from_secs(1), response_rx)
9315            .await
9316            .expect("a stalled ownership-bound send must not wedge the writer task")
9317            .unwrap();
9318        assert!(
9319            matches!(outcome, Err(SendError::WriteTimeout)),
9320            "expected the write deadline to be reported, was {outcome:?}"
9321        );
9322        assert_eq!(
9323            ConnectionMode::from_atomic(&connection_state),
9324            ConnectionMode::Reconnect
9325        );
9326
9327        // The timed-out message is ownership-bound, so unlike an ordinary send it must
9328        // NOT be buffered for replay onto the replacement sink.
9329        let recorded = Arc::new(Mutex::new(Vec::new()));
9330
9331        let recording_state = Arc::new(RecordingState {
9332            messages: Arc::clone(&recorded),
9333            recorded_notify: tokio::sync::Notify::new(),
9334        });
9335
9336        let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9337            state: Arc::clone(&recording_state),
9338        });
9339
9340        let (new_writer, _reader) = transport.split();
9341        let (update_tx, update_rx) = tokio::sync::oneshot::channel();
9342        writer_tx
9343            .send(WriterCommand::Update(new_writer, update_tx))
9344            .unwrap();
9345
9346        tokio::time::advance(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS)).await;
9347        assert_eq!(
9348            tokio::time::timeout(Duration::from_secs(1), update_rx)
9349                .await
9350                .expect("writer update should not remain queued behind a stalled send")
9351                .unwrap(),
9352            1
9353        );
9354
9355        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9356        state_notify.notify_waiters();
9357        tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9358
9359        // Barrier: two sentinels bound to the NEW epoch, the second sent only after the first
9360        // is acknowledged. The writer can accept the first from a `recv()` that was already
9361        // pending when the mode flipped, so only the second is necessarily processed after the
9362        // intervening loop-top drain - which is where a mistakenly buffered message would be
9363        // replayed. One sentinel alone would leave the ordering to chance.
9364        for name in ["sentinel-1", "sentinel-2"] {
9365            let (sentinel_tx, sentinel_rx) = tokio::sync::oneshot::channel();
9366            writer_tx
9367                .send_control(WriterCommand::SendOnConnection {
9368                    message: Message::text(name),
9369                    connection_epoch: 1,
9370                    response_tx: sentinel_tx,
9371                })
9372                .unwrap();
9373
9374            recording_state.recorded_notify.notified().await;
9375            sentinel_rx
9376                .await
9377                .unwrap()
9378                .expect("the sentinel should send on the replacement connection");
9379        }
9380
9381        assert_eq!(
9382            recorded.lock().as_slice(),
9383            &[Message::text("sentinel-1"), Message::text("sentinel-2")],
9384            "an ownership-bound message must never be replayed after its deadline expires"
9385        );
9386
9387        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9388        state_notify.notify_waiters();
9389        drop(writer_tx);
9390        write_task.await.unwrap();
9391    }
9392
9393    #[rstest]
9394    #[tokio::test(start_paused = true)]
9395    async fn test_stalled_websocket_replay_reconnects_and_retries_buffer() {
9396        let initial_messages = Arc::new(Mutex::new(Vec::new()));
9397
9398        let initial_recording_state = Arc::new(RecordingState {
9399            messages: initial_messages,
9400            recorded_notify: tokio::sync::Notify::new(),
9401        });
9402
9403        let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9404            state: initial_recording_state,
9405        });
9406
9407        let (writer, _reader) = transport.split();
9408        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
9409        let state_notify = Arc::new(tokio::sync::Notify::new());
9410        let auth_tracker = Arc::new(OnceLock::new());
9411        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
9412        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9413        let write_task = WebSocketClientInner::spawn_write_task(
9414            Arc::clone(&connection_state),
9415            Arc::clone(&state_notify),
9416            Arc::new(AtomicBool::new(true)),
9417            writer,
9418            writer_rx,
9419            Arc::new(AtomicU64::new(0)),
9420            Arc::clone(&auth_tracker),
9421            reconnect_buffer_waits_for_auth,
9422            None,
9423        );
9424
9425        writer_tx
9426            .send(WriterCommand::Send(Message::text("buffered-message")))
9427            .unwrap();
9428
9429        let blocking_state = Arc::new(BlockingFailState::default());
9430
9431        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9432            state: Arc::clone(&blocking_state),
9433        });
9434
9435        let (blocking_writer, _reader) = transport.split();
9436        let (blocking_tx, blocking_rx) = tokio::sync::oneshot::channel();
9437        writer_tx
9438            .send(WriterCommand::Update(blocking_writer, blocking_tx))
9439            .unwrap();
9440        assert_eq!(blocking_rx.await.unwrap(), 1);
9441
9442        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9443        state_notify.notify_waiters();
9444        tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9445        blocking_state.send_entered_notify.notified().await;
9446
9447        let recorded = Arc::new(Mutex::new(Vec::new()));
9448
9449        let recording_state = Arc::new(RecordingState {
9450            messages: Arc::clone(&recorded),
9451            recorded_notify: tokio::sync::Notify::new(),
9452        });
9453
9454        let transport: BoxedWsTransport = Box::pin(RecordingTransport {
9455            state: Arc::clone(&recording_state),
9456        });
9457
9458        let (new_writer, _reader) = transport.split();
9459        let (update_tx, update_rx) = tokio::sync::oneshot::channel();
9460        writer_tx
9461            .send(WriterCommand::Update(new_writer, update_tx))
9462            .unwrap();
9463
9464        tokio::time::advance(Duration::from_secs(GRACEFUL_SHUTDOWN_TIMEOUT_SECS)).await;
9465        assert_eq!(
9466            tokio::time::timeout(Duration::from_secs(1), update_rx)
9467                .await
9468                .expect("writer update should not remain queued behind stalled replay")
9469                .unwrap(),
9470            2,
9471            "the second replacement sink should install as connection epoch 2"
9472        );
9473        assert_eq!(
9474            ConnectionMode::from_atomic(&connection_state),
9475            ConnectionMode::Reconnect
9476        );
9477
9478        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9479        state_notify.notify_waiters();
9480        tokio::time::advance(Duration::from_millis(CONNECTION_STATE_CHECK_INTERVAL_MS)).await;
9481        recording_state.recorded_notify.notified().await;
9482        assert_eq!(
9483            recorded.lock().as_slice(),
9484            &[Message::text("buffered-message")]
9485        );
9486
9487        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9488        state_notify.notify_waiters();
9489        drop(writer_tx);
9490        write_task.await.unwrap();
9491    }
9492
9493    #[rstest]
9494    #[tokio::test]
9495    async fn test_new_with_writer_rejects_zero_heartbeat() {
9496        // Stream mode shares connect_url's validation: a zero heartbeat would
9497        // spawn a busy-loop ping flood
9498        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9499            state: Arc::new(BlockingFailState::default()),
9500        });
9501
9502        let (writer, _reader) = transport.split();
9503
9504        let config = WebSocketConfig {
9505            url: "ws://127.0.0.1:1".to_string(),
9506            headers: vec![],
9507            heartbeat_interval_secs: Some(0),
9508            heartbeat_payload: None,
9509            connect_timeout_ms: None,
9510            reconnect_delay_initial_ms: None,
9511            reconnect_delay_max_ms: None,
9512            reconnect_backoff_factor: None,
9513            reconnect_jitter_ms: None,
9514            reconnect_max_attempts: None,
9515            heartbeat_timeout_secs: None,
9516            idle_timeout_ms: None,
9517            writer_capacity: None,
9518            backend: TransportBackend::Tungstenite,
9519            proxy_url: None,
9520            max_message_size_bytes: None,
9521            max_frame_size_bytes: None,
9522        };
9523
9524        let err = WebSocketClientInner::new_with_writer(config, writer)
9525            .await
9526            .expect_err("zero heartbeat should be rejected in stream mode");
9527        assert!(
9528            err.to_string()
9529                .contains("Heartbeat interval cannot be zero"),
9530            "error should mention zero heartbeat, was: {err}"
9531        );
9532    }
9533
9534    #[rstest]
9535    #[tokio::test]
9536    async fn test_connect_times_out_on_silent_server() {
9537        // A server that accepts TCP but never completes the WebSocket upgrade
9538        // must not hang connect() indefinitely
9539        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
9540        let port = listener.local_addr().unwrap().port();
9541
9542        let server = task::spawn(async move {
9543            // Accept and hold the socket open without responding
9544            if let Ok((_stream, _)) = listener.accept().await {
9545                sleep(Duration::from_secs(30)).await;
9546            }
9547        });
9548
9549        let (handler, _rx) = channel_message_handler();
9550
9551        let config = WebSocketConfig {
9552            url: format!("ws://127.0.0.1:{port}"),
9553            headers: vec![],
9554            heartbeat_interval_secs: None,
9555            heartbeat_payload: None,
9556            connect_timeout_ms: Some(500),
9557            reconnect_delay_initial_ms: Some(50),
9558            reconnect_delay_max_ms: Some(100),
9559            reconnect_backoff_factor: Some(1.0),
9560            reconnect_jitter_ms: Some(0),
9561            reconnect_max_attempts: None,
9562            heartbeat_timeout_secs: None,
9563            idle_timeout_ms: None,
9564            writer_capacity: None,
9565            backend: TransportBackend::Tungstenite,
9566            proxy_url: None,
9567            max_message_size_bytes: None,
9568            max_frame_size_bytes: None,
9569        };
9570
9571        let result = tokio::time::timeout(
9572            Duration::from_secs(5),
9573            WebSocketClient::builder()
9574                .config(config)
9575                .message_handler(handler)
9576                .connect(),
9577        )
9578        .await
9579        .expect("connect should not hang on a silent server");
9580
9581        assert!(result.is_err(), "connect should fail with a timeout error");
9582        let err_msg = result.unwrap_err().to_string();
9583        assert!(
9584            err_msg.contains("timed out"),
9585            "error should mention the timeout, was: {err_msg}"
9586        );
9587
9588        server.abort();
9589    }
9590
9591    #[rstest]
9592    #[tokio::test]
9593    async fn test_reconnect_succeeds_with_timeout_shorter_than_swap_ceremony() {
9594        // Regression: the reconnect timeout used to cover the writer swap and
9595        // graceful-shutdown delays (~200ms minimum), so a short timeout caused
9596        // every otherwise-successful reconnect to be discarded mid-swap and the
9597        // already-swapped writer to be orphaned. The timeout now bounds only
9598        // connection establishment.
9599        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
9600        let port = listener.local_addr().unwrap().port();
9601        let (release_first_tx, release_first_rx) = tokio::sync::oneshot::channel();
9602
9603        let server = task::spawn(async move {
9604            // Keep the first socket until connect() returns; an immediate drop
9605            // can fail the initial handshake
9606            if let Ok((stream, _)) = listener.accept().await
9607                && let Ok(ws) = accept_async(stream).await
9608            {
9609                let _ = release_first_rx.await;
9610                drop(ws);
9611            }
9612
9613            // A timed-out reconnect can consume an accept without completing the
9614            // handshake, so keep accepting and announce on every replacement
9615            loop {
9616                let Ok((stream, _)) = listener.accept().await else {
9617                    break;
9618                };
9619
9620                if let Ok(mut ws) = accept_async(stream).await {
9621                    loop {
9622                        if ws
9623                            .send(WsMessage::Text("reconnected-msg".to_string().into()))
9624                            .await
9625                            .is_err()
9626                        {
9627                            break;
9628                        }
9629
9630                        sleep(Duration::from_millis(50)).await;
9631                    }
9632                }
9633            }
9634        });
9635
9636        let (handler, mut rx) = channel_message_handler();
9637
9638        let config = WebSocketConfig {
9639            url: format!("ws://127.0.0.1:{port}"),
9640            headers: vec![],
9641            heartbeat_interval_secs: None,
9642            heartbeat_payload: None,
9643            connect_timeout_ms: Some(150), // Shorter than the ~200ms swap ceremony
9644            reconnect_delay_initial_ms: Some(25),
9645            reconnect_delay_max_ms: Some(50),
9646            reconnect_backoff_factor: Some(1.0),
9647            reconnect_jitter_ms: Some(0),
9648            reconnect_max_attempts: None,
9649            heartbeat_timeout_secs: None,
9650            idle_timeout_ms: None,
9651            writer_capacity: None,
9652            backend: TransportBackend::Tungstenite,
9653            proxy_url: None,
9654            max_message_size_bytes: None,
9655            max_frame_size_bytes: None,
9656        };
9657
9658        let client = WebSocketClient::builder()
9659            .config(config)
9660            .message_handler(handler)
9661            .connect()
9662            .await
9663            .expect("Initial WebSocket connection should succeed before triggering reconnect");
9664        wait_until_async(|| async { client.is_active() }, Duration::from_secs(2)).await;
9665        release_first_tx
9666            .send(())
9667            .expect("server should still be holding the first connection");
9668
9669        let received = tokio::time::timeout(Duration::from_secs(5), async {
9670            loop {
9671                match rx.recv().await {
9672                    Some(WsMessage::Text(text)) if text.as_str() == "reconnected-msg" => {
9673                        return true;
9674                    }
9675                    Some(_) => {}
9676                    None => return false,
9677                }
9678            }
9679        })
9680        .await
9681        .expect("Timed out waiting for a message from the replacement WebSocket connection");
9682
9683        assert!(
9684            received,
9685            "Replacement WebSocket message channel closed before delivering the reconnect message"
9686        );
9687        assert!(
9688            client.is_active(),
9689            "WebSocket client should be active after delivering the reconnect message"
9690        );
9691
9692        client.disconnect().await;
9693        server.abort();
9694    }
9695
9696    #[rstest]
9697    #[tokio::test]
9698    async fn test_idle_timeout_fires_under_ping_flood() {
9699        // Regression: the idle check used to run only when nothing arrived for a
9700        // full check interval (10ms), so pings flooding faster than that starved
9701        // it and a ping-only zombie connection never tripped the timeout
9702        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
9703        let port = listener.local_addr().unwrap().port();
9704        let server = task::spawn(async move {
9705            let (stream, _) = listener.accept().await.unwrap();
9706            let mut ws = accept_async(stream).await.unwrap();
9707
9708            for _ in 0..600 {
9709                sleep(Duration::from_millis(5)).await;
9710
9711                if ws.send(WsMessage::Ping(Vec::new().into())).await.is_err() {
9712                    break;
9713                }
9714            }
9715        });
9716
9717        let (handler, _rx) = channel_message_handler();
9718
9719        let config = WebSocketConfig {
9720            url: format!("ws://127.0.0.1:{port}"),
9721            headers: vec![],
9722            heartbeat_interval_secs: None,
9723            heartbeat_payload: None,
9724            connect_timeout_ms: Some(2_000),
9725            reconnect_delay_initial_ms: Some(50),
9726            reconnect_delay_max_ms: Some(100),
9727            reconnect_backoff_factor: Some(1.0),
9728            reconnect_jitter_ms: Some(0),
9729            reconnect_max_attempts: Some(1),
9730            heartbeat_timeout_secs: None,
9731            idle_timeout_ms: Some(500),
9732            writer_capacity: None,
9733            backend: TransportBackend::Tungstenite,
9734            proxy_url: None,
9735            max_message_size_bytes: None,
9736            max_frame_size_bytes: None,
9737        };
9738
9739        let client = WebSocketClient::builder()
9740            .config(config)
9741            .message_handler(handler)
9742            .connect()
9743            .await
9744            .unwrap();
9745
9746        assert!(client.is_active());
9747
9748        wait_until_async(
9749            || async { client.is_reconnecting() || client.is_disconnected() },
9750            Duration::from_millis(1_500),
9751        )
9752        .await;
9753
9754        assert!(
9755            !client.is_active(),
9756            "Client should not be active after idle timeout under a ping flood"
9757        );
9758
9759        client.disconnect().await;
9760        server.abort();
9761    }
9762
9763    #[rstest]
9764    #[tokio::test]
9765    async fn test_send_failure_does_not_overwrite_disconnect() {
9766        let state = Arc::new(BlockingFailState::default());
9767
9768        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9769            state: Arc::clone(&state),
9770        });
9771
9772        let (writer, _reader) = transport.split();
9773
9774        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9775        let state_notify = Arc::new(tokio::sync::Notify::new());
9776        let auth_tracker = Arc::new(OnceLock::new());
9777        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(false));
9778        let connection_epoch = Arc::new(AtomicU64::new(0));
9779
9780        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9781        let write_task = WebSocketClientInner::spawn_write_task(
9782            Arc::clone(&connection_state),
9783            Arc::clone(&state_notify),
9784            Arc::new(AtomicBool::new(true)),
9785            writer,
9786            writer_rx,
9787            connection_epoch,
9788            Arc::clone(&auth_tracker),
9789            Arc::clone(&reconnect_buffer_waits_for_auth),
9790            None,
9791        );
9792
9793        writer_tx
9794            .send(WriterCommand::Send(Message::text("doomed")))
9795            .unwrap();
9796
9797        // Wait until the writer task is blocked inside the transport send
9798        wait_until_async(
9799            || {
9800                let state = Arc::clone(&state);
9801                async move { state.send_entered.load(Ordering::SeqCst) }
9802            },
9803            Duration::from_secs(2),
9804        )
9805        .await;
9806
9807        // Disconnect lands while the send is in flight, then the send fails;
9808        // the writer error path must not overwrite DISCONNECT with RECONNECT
9809        connection_state.store(ConnectionMode::Disconnect.as_u8(), Ordering::SeqCst);
9810        state.trigger_failure();
9811
9812        tokio::time::timeout(Duration::from_secs(2), write_task)
9813            .await
9814            .expect("write task should exit after disconnect")
9815            .unwrap();
9816
9817        assert_eq!(
9818            ConnectionMode::from_atomic(&connection_state),
9819            ConnectionMode::Disconnect,
9820            "Send failure must not resurrect a disconnecting client into RECONNECT"
9821        );
9822    }
9823
9824    #[tokio::test]
9825    async fn send_on_connection_write_failure_reports_broken_pipe_and_reconnects() {
9826        let state = Arc::new(BlockingFailState::default());
9827
9828        let transport: BoxedWsTransport = Box::pin(BlockingFailTransport {
9829            state: Arc::clone(&state),
9830        });
9831
9832        let (writer, _reader) = transport.split();
9833        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9834        let state_notify = Arc::new(tokio::sync::Notify::new());
9835        let connection_epoch = Arc::new(AtomicU64::new(0));
9836        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9837        let write_task = WebSocketClientInner::spawn_write_task(
9838            Arc::clone(&connection_state),
9839            Arc::clone(&state_notify),
9840            Arc::new(AtomicBool::new(true)),
9841            writer,
9842            writer_rx,
9843            Arc::clone(&connection_epoch),
9844            Arc::new(OnceLock::new()),
9845            Arc::new(AtomicBool::new(false)),
9846            None,
9847        );
9848
9849        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9850        writer_tx
9851            .send_control(WriterCommand::SendOnConnection {
9852                message: Message::text("doomed"),
9853                connection_epoch: 0,
9854                response_tx,
9855            })
9856            .unwrap();
9857
9858        wait_until_async(
9859            || {
9860                let state = Arc::clone(&state);
9861                async move { state.send_entered.load(Ordering::SeqCst) }
9862            },
9863            Duration::from_secs(2),
9864        )
9865        .await;
9866
9867        state.trigger_failure();
9868
9869        match response_rx.await.unwrap().unwrap_err() {
9870            SendError::BrokenPipe(message) => assert_eq!(message, "connection reset"),
9871            other => panic!("expected broken-pipe send error, was {other:?}"),
9872        }
9873
9874        wait_until_async(
9875            || async {
9876                ConnectionMode::from_atomic(&connection_state) == ConnectionMode::Reconnect
9877            },
9878            Duration::from_secs(2),
9879        )
9880        .await;
9881
9882        assert_eq!(
9883            ConnectionMode::from_atomic(&connection_state),
9884            ConnectionMode::Reconnect,
9885        );
9886        assert_eq!(connection_epoch.load(Ordering::Acquire), 0);
9887
9888        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
9889        state_notify.notify_waiters();
9890        drop(writer_tx);
9891        write_task.await.unwrap();
9892    }
9893
9894    #[tokio::test]
9895    async fn send_on_connection_rejects_stale_epoch_without_replay() {
9896        let server = RecordingServer::setup().await;
9897        let url = format!("ws://127.0.0.1:{}", server.port);
9898        let (writer, _reader) = WebSocketClientInner::connect_with_server(
9899            &url,
9900            vec![],
9901            TransportBackend::Tungstenite,
9902            None,
9903            None,
9904            None,
9905        )
9906        .await
9907        .unwrap();
9908        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Active.as_u8()));
9909        let state_notify = Arc::new(tokio::sync::Notify::new());
9910        let auth_tracker = Arc::new(OnceLock::new());
9911        let connection_epoch = Arc::new(AtomicU64::new(0));
9912        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
9913        let write_task = WebSocketClientInner::spawn_write_task(
9914            Arc::clone(&connection_state),
9915            Arc::clone(&state_notify),
9916            Arc::new(AtomicBool::new(true)),
9917            writer,
9918            writer_rx,
9919            Arc::clone(&connection_epoch),
9920            auth_tracker,
9921            Arc::new(AtomicBool::new(false)),
9922            None,
9923        );
9924
9925        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9926        writer_tx
9927            .send_control(WriterCommand::SendOnConnection {
9928                message: Message::text("epoch-0"),
9929                connection_epoch: 0,
9930                response_tx,
9931            })
9932            .unwrap();
9933
9934        response_rx.await.unwrap().unwrap();
9935
9936        connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
9937        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9938        writer_tx
9939            .send_control(WriterCommand::SendOnConnection {
9940                message: Message::text("during-reconnect"),
9941                connection_epoch: 0,
9942                response_tx,
9943            })
9944            .unwrap();
9945
9946        assert!(matches!(
9947            response_rx.await.unwrap(),
9948            Err(SendError::ConnectionChanged),
9949        ));
9950
9951        let (replacement, _reader) = WebSocketClientInner::connect_with_server(
9952            &url,
9953            vec![],
9954            TransportBackend::Tungstenite,
9955            None,
9956            None,
9957            None,
9958        )
9959        .await
9960        .unwrap();
9961        let (update_tx, update_rx) = tokio::sync::oneshot::channel();
9962        writer_tx
9963            .send(WriterCommand::Update(replacement, update_tx))
9964            .unwrap();
9965        assert_eq!(update_rx.await.unwrap(), 1);
9966        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
9967
9968        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9969        writer_tx
9970            .send_control(WriterCommand::SendOnConnection {
9971                message: Message::text("stale"),
9972                connection_epoch: 0,
9973                response_tx,
9974            })
9975            .unwrap();
9976
9977        assert!(matches!(
9978            response_rx.await.unwrap(),
9979            Err(SendError::ConnectionChanged),
9980        ));
9981
9982        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
9983        writer_tx
9984            .send_control(WriterCommand::SendOnConnection {
9985                message: Message::text("epoch-1"),
9986                connection_epoch: 1,
9987                response_tx,
9988            })
9989            .unwrap();
9990
9991        response_rx.await.unwrap().unwrap();
9992
9993        connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
9994        let (second_replacement, _reader) = WebSocketClientInner::connect_with_server(
9995            &url,
9996            vec![],
9997            TransportBackend::Tungstenite,
9998            None,
9999            None,
10000            None,
10001        )
10002        .await
10003        .unwrap();
10004        let (update_tx, update_rx) = tokio::sync::oneshot::channel();
10005        writer_tx
10006            .send(WriterCommand::Update(second_replacement, update_tx))
10007            .unwrap();
10008        assert_eq!(update_rx.await.unwrap(), 2);
10009        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10010
10011        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10012        writer_tx
10013            .send_control(WriterCommand::SendOnConnection {
10014                message: Message::text("stale-after-second-reconnect"),
10015                connection_epoch: 1,
10016                response_tx,
10017            })
10018            .unwrap();
10019
10020        assert!(matches!(
10021            response_rx.await.unwrap(),
10022            Err(SendError::ConnectionChanged),
10023        ));
10024
10025        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10026        writer_tx
10027            .send_control(WriterCommand::SendOnConnection {
10028                message: Message::text("epoch-2"),
10029                connection_epoch: 2,
10030                response_tx,
10031            })
10032            .unwrap();
10033
10034        response_rx.await.unwrap().unwrap();
10035
10036        wait_until_async(
10037            || {
10038                let messages = Arc::clone(&server.messages);
10039                async move { messages.lock().await.len() == 3 }
10040            },
10041            Duration::from_secs(2),
10042        )
10043        .await;
10044
10045        assert_eq!(connection_epoch.load(Ordering::Acquire), 2);
10046        assert_eq!(
10047            server.messages().await,
10048            vec![
10049                "epoch-0".to_string(),
10050                "epoch-1".to_string(),
10051                "epoch-2".to_string(),
10052            ],
10053        );
10054
10055        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10056        state_notify.notify_waiters();
10057        drop(writer_tx);
10058        write_task.abort();
10059    }
10060
10061    #[rstest]
10062    fn test_reconnect_buffer_action_requires_active_mode() {
10063        let connection_state = AtomicU8::new(ConnectionMode::Active.as_u8());
10064        let reconnect_buffer_waits_for_auth = AtomicBool::new(true);
10065        let auth_tracker = Arc::new(OnceLock::new());
10066        let tracker = AuthTracker::new();
10067        tracker.succeed();
10068        auth_tracker.set(tracker).unwrap();
10069
10070        assert_eq!(
10071            WebSocketClientInner::reconnect_buffer_action(
10072                &reconnect_buffer_waits_for_auth,
10073                &auth_tracker,
10074                &connection_state,
10075            ),
10076            ReconnectBufferAction::Drain,
10077        );
10078
10079        connection_state.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
10080        assert_eq!(
10081            WebSocketClientInner::reconnect_buffer_action(
10082                &reconnect_buffer_waits_for_auth,
10083                &auth_tracker,
10084                &connection_state,
10085            ),
10086            ReconnectBufferAction::Wait,
10087        );
10088    }
10089
10090    #[tokio::test]
10091    async fn test_write_task_waits_for_auth_before_replaying_buffer() {
10092        use nautilus_common::testing::wait_until_async;
10093
10094        let server = RecordingServer::setup().await;
10095        let url = format!("ws://127.0.0.1:{}", server.port);
10096        let (writer, _reader) = WebSocketClientInner::connect_with_server(
10097            &url,
10098            vec![],
10099            TransportBackend::Tungstenite,
10100            None,
10101            None,
10102            None,
10103        )
10104        .await
10105        .unwrap();
10106
10107        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
10108        let state_notify = Arc::new(tokio::sync::Notify::new());
10109        let auth_tracker = Arc::new(OnceLock::new());
10110        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(true));
10111        let connection_epoch = Arc::new(AtomicU64::new(0));
10112        let tracker = AuthTracker::new();
10113        auth_tracker.set(tracker.clone()).unwrap();
10114
10115        let (writer_tx, writer_rx) = crate::writer::channel(1);
10116        let write_task = WebSocketClientInner::spawn_write_task(
10117            Arc::clone(&connection_state),
10118            Arc::clone(&state_notify),
10119            Arc::new(AtomicBool::new(true)),
10120            writer,
10121            writer_rx,
10122            Arc::clone(&connection_epoch),
10123            Arc::clone(&auth_tracker),
10124            Arc::clone(&reconnect_buffer_waits_for_auth),
10125            None,
10126        );
10127
10128        writer_tx
10129            .send(WriterCommand::Send(Message::Text("stale".into())))
10130            .unwrap();
10131
10132        let (new_writer, _reader) = WebSocketClientInner::connect_with_server(
10133            &url,
10134            vec![],
10135            TransportBackend::Tungstenite,
10136            None,
10137            None,
10138            None,
10139        )
10140        .await
10141        .unwrap();
10142        let (tx, rx) = tokio::sync::oneshot::channel();
10143        writer_tx
10144            .send_update(WriterCommand::Update(new_writer, tx))
10145            .unwrap();
10146        assert_eq!(rx.await.unwrap(), 1);
10147
10148        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10149
10150        tokio::time::sleep(Duration::from_millis(300)).await;
10151        assert!(
10152            server.messages().await.is_empty(),
10153            "buffered messages should wait for re-authentication"
10154        );
10155
10156        assert!(matches!(
10157            writer_tx.send(WriterCommand::Send(Message::text("overflow"))),
10158            Err(SendError::BufferFull)
10159        ));
10160        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10161        writer_tx
10162            .send_control(WriterCommand::SendOnConnection {
10163                message: Message::text("authenticate"),
10164                connection_epoch: 1,
10165                response_tx,
10166            })
10167            .unwrap();
10168
10169        response_rx.await.unwrap().unwrap();
10170        tracker.succeed();
10171
10172        wait_until_async(
10173            || {
10174                let messages = Arc::clone(&server.messages);
10175                async move { messages.lock().await.len() == 2 }
10176            },
10177            Duration::from_secs(3),
10178        )
10179        .await;
10180
10181        assert_eq!(
10182            server.messages().await,
10183            vec!["authenticate".to_string(), "stale".to_string()]
10184        );
10185
10186        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10187        state_notify.notify_waiters();
10188        drop(writer_tx);
10189        write_task.abort();
10190    }
10191
10192    #[rstest]
10193    #[tokio::test]
10194    async fn test_write_task_buffers_send_behind_pending_replay() {
10195        use nautilus_common::testing::wait_until_async;
10196
10197        let server = RecordingServer::setup().await;
10198        let url = format!("ws://127.0.0.1:{}", server.port);
10199        let (writer, _reader) = WebSocketClientInner::connect_with_server(
10200            &url,
10201            vec![],
10202            TransportBackend::Tungstenite,
10203            None,
10204            None,
10205            None,
10206        )
10207        .await
10208        .unwrap();
10209
10210        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
10211        let state_notify = Arc::new(tokio::sync::Notify::new());
10212        let auth_tracker = Arc::new(OnceLock::new());
10213        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(true));
10214        let connection_epoch = Arc::new(AtomicU64::new(0));
10215        let tracker = AuthTracker::new();
10216        auth_tracker.set(tracker.clone()).unwrap();
10217
10218        // Two slots: the buffered message holds one while replay waits for auth, so a second
10219        // ordinary send is still admissible and reaches the writer loop.
10220        let (writer_tx, writer_rx) = crate::writer::channel(2);
10221        let write_task = WebSocketClientInner::spawn_write_task(
10222            Arc::clone(&connection_state),
10223            Arc::clone(&state_notify),
10224            Arc::new(AtomicBool::new(true)),
10225            writer,
10226            writer_rx,
10227            Arc::clone(&connection_epoch),
10228            Arc::clone(&auth_tracker),
10229            Arc::clone(&reconnect_buffer_waits_for_auth),
10230            None,
10231        );
10232
10233        writer_tx
10234            .send(WriterCommand::Send(Message::Text("stale".into())))
10235            .unwrap();
10236
10237        let (new_writer, _reader) = WebSocketClientInner::connect_with_server(
10238            &url,
10239            vec![],
10240            TransportBackend::Tungstenite,
10241            None,
10242            None,
10243            None,
10244        )
10245        .await
10246        .unwrap();
10247        let (tx, rx) = tokio::sync::oneshot::channel();
10248        writer_tx
10249            .send_update(WriterCommand::Update(new_writer, tx))
10250            .unwrap();
10251        assert_eq!(rx.await.unwrap(), 1);
10252
10253        // Active with a non-empty replay buffer: the drain is gated on authentication.
10254        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10255        tokio::time::sleep(Duration::from_millis(300)).await;
10256        assert!(
10257            server.messages().await.is_empty(),
10258            "buffered message should wait for re-authentication"
10259        );
10260
10261        // A later ordinary send must not overtake the still-buffered message, nor bypass
10262        // the authentication gate by writing to the replacement connection.
10263        writer_tx
10264            .send(WriterCommand::Send(Message::Text("fresh".into())))
10265            .unwrap();
10266        tokio::time::sleep(Duration::from_millis(300)).await;
10267        assert!(
10268            server.messages().await.is_empty(),
10269            "a send arriving before the gate opens must not reach the connection, was {:?}",
10270            server.messages().await
10271        );
10272
10273        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10274        writer_tx
10275            .send_control(WriterCommand::SendOnConnection {
10276                message: Message::text("authenticate"),
10277                connection_epoch: 1,
10278                response_tx,
10279            })
10280            .unwrap();
10281
10282        response_rx.await.unwrap().unwrap();
10283        tracker.succeed();
10284
10285        wait_until_async(
10286            || {
10287                let messages = Arc::clone(&server.messages);
10288                async move { messages.lock().await.len() == 3 }
10289            },
10290            Duration::from_secs(3),
10291        )
10292        .await;
10293
10294        assert_eq!(
10295            server.messages().await,
10296            vec![
10297                "authenticate".to_string(),
10298                "stale".to_string(),
10299                "fresh".to_string(),
10300            ],
10301            "replay must preserve FIFO order behind the authentication handshake"
10302        );
10303
10304        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10305        state_notify.notify_waiters();
10306        drop(writer_tx);
10307        write_task.abort();
10308    }
10309
10310    #[rstest]
10311    #[tokio::test]
10312    async fn test_write_task_control_send_reaches_connection_behind_pending_replay() {
10313        use nautilus_common::testing::wait_until_async;
10314
10315        let server = RecordingServer::setup().await;
10316        let url = format!("ws://127.0.0.1:{}", server.port);
10317        let (writer, _reader) = WebSocketClientInner::connect_with_server(
10318            &url,
10319            vec![],
10320            TransportBackend::Tungstenite,
10321            None,
10322            None,
10323            None,
10324        )
10325        .await
10326        .unwrap();
10327
10328        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
10329        let state_notify = Arc::new(tokio::sync::Notify::new());
10330        let auth_tracker = Arc::new(OnceLock::new());
10331        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(true));
10332        let connection_epoch = Arc::new(AtomicU64::new(0));
10333        let tracker = AuthTracker::new();
10334        auth_tracker.set(tracker.clone()).unwrap();
10335
10336        let (writer_tx, writer_rx) = crate::writer::channel(2);
10337        let write_task = WebSocketClientInner::spawn_write_task(
10338            Arc::clone(&connection_state),
10339            Arc::clone(&state_notify),
10340            Arc::new(AtomicBool::new(true)),
10341            writer,
10342            writer_rx,
10343            Arc::clone(&connection_epoch),
10344            Arc::clone(&auth_tracker),
10345            Arc::clone(&reconnect_buffer_waits_for_auth),
10346            None,
10347        );
10348
10349        writer_tx
10350            .send(WriterCommand::Send(Message::Text("stale".into())))
10351            .unwrap();
10352
10353        let (new_writer, _reader) = WebSocketClientInner::connect_with_server(
10354            &url,
10355            vec![],
10356            TransportBackend::Tungstenite,
10357            None,
10358            None,
10359            None,
10360        )
10361        .await
10362        .unwrap();
10363        let (tx, rx) = tokio::sync::oneshot::channel();
10364        writer_tx
10365            .send_update(WriterCommand::Update(new_writer, tx))
10366            .unwrap();
10367        assert_eq!(rx.await.unwrap(), 1);
10368
10369        // Active with a non-empty replay buffer: the drain is gated on authentication.
10370        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10371        tokio::time::sleep(Duration::from_millis(300)).await;
10372        assert!(
10373            server.messages().await.is_empty(),
10374            "buffered message should wait for re-authentication"
10375        );
10376
10377        // A control frame enqueued through the ordinary send path in the same window must not be
10378        // buffered with the replay: it has to reach the replacement connection while the gate is
10379        // still closed.
10380        writer_tx
10381            .send_control(WriterCommand::Send(Message::Pong(
10382                b"keepalive".to_vec().into(),
10383            )))
10384            .unwrap();
10385
10386        wait_until_async(
10387            || {
10388                let control_frames = Arc::clone(&server.control_frames);
10389                async move { !control_frames.lock().await.is_empty() }
10390            },
10391            Duration::from_secs(3),
10392        )
10393        .await;
10394
10395        assert_eq!(
10396            server.control_frames().await,
10397            vec!["pong:keepalive".to_string()],
10398            "a control frame sent during the wait must reach the new connection without replay"
10399        );
10400        assert!(
10401            server.messages().await.is_empty(),
10402            "ordinary sends must still wait for re-authentication, was {:?}",
10403            server.messages().await
10404        );
10405
10406        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
10407        writer_tx
10408            .send_control(WriterCommand::SendOnConnection {
10409                message: Message::text("authenticate"),
10410                connection_epoch: 1,
10411                response_tx,
10412            })
10413            .unwrap();
10414
10415        response_rx.await.unwrap().unwrap();
10416        tracker.succeed();
10417
10418        wait_until_async(
10419            || {
10420                let messages = Arc::clone(&server.messages);
10421                async move { messages.lock().await.len() == 2 }
10422            },
10423            Duration::from_secs(3),
10424        )
10425        .await;
10426
10427        assert_eq!(
10428            server.messages().await,
10429            vec!["authenticate".to_string(), "stale".to_string()],
10430            "the control frame must not disturb ordinary replay ordering"
10431        );
10432
10433        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10434        state_notify.notify_waiters();
10435        drop(writer_tx);
10436        write_task.abort();
10437    }
10438
10439    #[tokio::test]
10440    async fn test_write_task_discards_buffer_after_auth_failure() {
10441        let server = RecordingServer::setup().await;
10442        let url = format!("ws://127.0.0.1:{}", server.port);
10443        let (writer, _reader) = WebSocketClientInner::connect_with_server(
10444            &url,
10445            vec![],
10446            TransportBackend::Tungstenite,
10447            None,
10448            None,
10449            None,
10450        )
10451        .await
10452        .unwrap();
10453
10454        let connection_state = Arc::new(AtomicU8::new(ConnectionMode::Reconnect.as_u8()));
10455        let state_notify = Arc::new(tokio::sync::Notify::new());
10456        let auth_tracker = Arc::new(OnceLock::new());
10457        let reconnect_buffer_waits_for_auth = Arc::new(AtomicBool::new(true));
10458        let connection_epoch = Arc::new(AtomicU64::new(0));
10459        let tracker = AuthTracker::new();
10460        auth_tracker.set(tracker.clone()).unwrap();
10461
10462        let (writer_tx, writer_rx) = crate::writer::channel(crate::writer::DEFAULT_WRITER_CAPACITY);
10463        let write_task = WebSocketClientInner::spawn_write_task(
10464            Arc::clone(&connection_state),
10465            Arc::clone(&state_notify),
10466            Arc::new(AtomicBool::new(true)),
10467            writer,
10468            writer_rx,
10469            Arc::clone(&connection_epoch),
10470            Arc::clone(&auth_tracker),
10471            Arc::clone(&reconnect_buffer_waits_for_auth),
10472            None,
10473        );
10474
10475        writer_tx
10476            .send(WriterCommand::Send(Message::Text("stale".into())))
10477            .unwrap();
10478
10479        let (new_writer, _reader) = WebSocketClientInner::connect_with_server(
10480            &url,
10481            vec![],
10482            TransportBackend::Tungstenite,
10483            None,
10484            None,
10485            None,
10486        )
10487        .await
10488        .unwrap();
10489        let (tx, rx) = tokio::sync::oneshot::channel();
10490        writer_tx
10491            .send(WriterCommand::Update(new_writer, tx))
10492            .unwrap();
10493        assert_eq!(rx.await.unwrap(), 1);
10494
10495        connection_state.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
10496        tracker.fail("rejected");
10497        tracker.invalidate();
10498        assert_eq!(tracker.auth_state(), AuthState::Failed);
10499        tokio::time::sleep(Duration::from_millis(300)).await;
10500        assert!(
10501            server.messages().await.is_empty(),
10502            "buffered messages should be discarded after authentication failure"
10503        );
10504
10505        let _auth_receiver = tracker.begin();
10506        tracker.succeed();
10507        tokio::time::sleep(Duration::from_millis(300)).await;
10508        assert!(
10509            server.messages().await.is_empty(),
10510            "discarded buffered messages should not replay on a later auth success"
10511        );
10512
10513        connection_state.store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
10514        state_notify.notify_waiters();
10515        drop(writer_tx);
10516        write_task.abort();
10517    }
10518
10519    #[rstest]
10520    #[tokio::test]
10521    async fn test_zero_idle_timeout_rejected() {
10522        let (handler, _rx) = channel_message_handler();
10523
10524        let config = WebSocketConfig {
10525            url: "ws://127.0.0.1:9999".to_string(),
10526            headers: vec![],
10527            heartbeat_interval_secs: None,
10528            heartbeat_payload: None,
10529            connect_timeout_ms: None,
10530            reconnect_delay_initial_ms: None,
10531            reconnect_delay_max_ms: None,
10532            reconnect_backoff_factor: None,
10533            reconnect_jitter_ms: None,
10534            reconnect_max_attempts: None,
10535            heartbeat_timeout_secs: None,
10536            idle_timeout_ms: Some(0),
10537            writer_capacity: None,
10538            backend: TransportBackend::Tungstenite,
10539            proxy_url: None,
10540            max_message_size_bytes: None,
10541            max_frame_size_bytes: None,
10542        };
10543
10544        let result = WebSocketClient::builder()
10545            .config(config)
10546            .message_handler(handler)
10547            .connect()
10548            .await;
10549
10550        assert!(result.is_err(), "Zero idle timeout should be rejected");
10551        let err_msg = result.unwrap_err().to_string();
10552        assert!(
10553            err_msg.contains("idle_timeout_ms"),
10554            "Error should name the offending field, was: {err_msg}"
10555        );
10556    }
10557
10558    #[rstest]
10559    #[tokio::test]
10560    async fn test_zero_heartbeat_timeout_rejected() {
10561        let (handler, _rx) = channel_message_handler();
10562
10563        let config = WebSocketConfig {
10564            url: "ws://127.0.0.1:9999".to_string(),
10565            headers: vec![],
10566            heartbeat_interval_secs: Some(30),
10567            heartbeat_payload: None,
10568            connect_timeout_ms: None,
10569            reconnect_delay_initial_ms: None,
10570            reconnect_delay_max_ms: None,
10571            reconnect_backoff_factor: None,
10572            reconnect_jitter_ms: None,
10573            reconnect_max_attempts: None,
10574            heartbeat_timeout_secs: Some(0),
10575            idle_timeout_ms: None,
10576            writer_capacity: None,
10577            backend: TransportBackend::Tungstenite,
10578            proxy_url: None,
10579            max_message_size_bytes: None,
10580            max_frame_size_bytes: None,
10581        };
10582
10583        let result = WebSocketClient::builder()
10584            .config(config)
10585            .message_handler(handler)
10586            .connect()
10587            .await;
10588
10589        assert!(result.is_err(), "Zero heartbeat timeout should be rejected");
10590        let err_msg = result.unwrap_err().to_string();
10591        assert!(
10592            err_msg.contains("heartbeat_timeout_secs"),
10593            "Error should name the offending field, was: {err_msg}"
10594        );
10595    }
10596
10597    #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
10598    #[rstest]
10599    #[tokio::test]
10600    async fn test_sockudo_backend_rejects_reserved_headers_before_connect() {
10601        let (handler, _rx) = channel_message_handler();
10602
10603        let config = WebSocketConfig {
10604            url: "ws://127.0.0.1:1".to_string(),
10605            headers: vec![("Host".to_string(), "example.com".to_string())],
10606            heartbeat_interval_secs: None,
10607            heartbeat_payload: None,
10608            connect_timeout_ms: None,
10609            reconnect_delay_initial_ms: None,
10610            reconnect_delay_max_ms: None,
10611            reconnect_backoff_factor: None,
10612            reconnect_jitter_ms: None,
10613            reconnect_max_attempts: None,
10614            heartbeat_timeout_secs: None,
10615            idle_timeout_ms: None,
10616            writer_capacity: None,
10617            backend: TransportBackend::Sockudo,
10618            proxy_url: None,
10619            max_message_size_bytes: None,
10620            max_frame_size_bytes: None,
10621        };
10622
10623        let err = WebSocketClient::builder()
10624            .config(config)
10625            .message_handler(handler)
10626            .connect()
10627            .await
10628            .expect_err("reserved header should fail before TCP connect");
10629
10630        assert!(
10631            err.to_string()
10632                .contains("reserved upgrade header not allowed in extra_headers"),
10633            "expected reserved-header failure, was: {err}"
10634        );
10635    }
10636
10637    // The log-capture harness is Linux-only for CI stability.
10638    #[cfg(all(feature = "transport-sockudo", target_os = "linux"))]
10639    #[rstest]
10640    #[tokio::test]
10641    async fn test_sockudo_reconnect_recovers_from_retryable_upgrade_rejection() {
10642        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
10643        let port = listener.local_addr().unwrap().port();
10644
10645        let server = task::spawn(async move {
10646            // First connection: accept the upgrade, then drop it to force a reconnect.
10647            let (stream, _) = listener.accept().await.unwrap();
10648            drop(accept_async(stream).await.unwrap());
10649
10650            // Second connection: reject the reconnect handshake with a retryable status.
10651            let (mut stream, _) = listener.accept().await.unwrap();
10652            let _request = read_http_request(&mut stream).await;
10653            stream
10654                .write_all(b"HTTP/1.1 502 Bad Gateway\r\n\r\n")
10655                .await
10656                .unwrap();
10657            stream.flush().await.unwrap();
10658            sleep(Duration::from_millis(50)).await;
10659            drop(stream);
10660
10661            // Third connection: accept the upgrade so the reconnect loop completes recovery.
10662            let (stream, _) = listener.accept().await.unwrap();
10663            let mut ws = accept_async(stream).await.unwrap();
10664            ws.send(WsMessage::Text("recovered".to_string().into()))
10665                .await
10666                .unwrap();
10667            sleep(Duration::from_secs(1)).await;
10668        });
10669
10670        let (handler, mut rx) = channel_message_handler();
10671
10672        let config = WebSocketConfig {
10673            url: format!(
10674                "ws://user-secret:password-secret@127.0.0.1:{port}/path-secret?token=query-secret"
10675            ),
10676            headers: vec![],
10677            heartbeat_interval_secs: None,
10678            heartbeat_payload: None,
10679            connect_timeout_ms: Some(2_000),
10680            reconnect_delay_initial_ms: Some(50),
10681            reconnect_delay_max_ms: Some(100),
10682            reconnect_backoff_factor: Some(1.0),
10683            reconnect_jitter_ms: Some(0),
10684            reconnect_max_attempts: None,
10685            heartbeat_timeout_secs: None,
10686            idle_timeout_ms: None,
10687            writer_capacity: None,
10688            backend: TransportBackend::Sockudo,
10689            proxy_url: None,
10690            max_message_size_bytes: None,
10691            max_frame_size_bytes: None,
10692        };
10693
10694        // Capture before connecting so the whole reconnect cycle is observed. The harness owns
10695        // the global logger, so the Nautilus logger that arms `shutdown_on_error` on an ERROR
10696        // record cannot run here; the absence of ERROR records is what rules the trigger out.
10697        let capture = capture_logs_for(RECONNECT_LOG_TARGETS).await;
10698
10699        let client = WebSocketClient::builder()
10700            .config(config)
10701            .message_handler(handler)
10702            .connect()
10703            .await
10704            .expect("sockudo connect should succeed");
10705
10706        let recovered = tokio::time::timeout(Duration::from_secs(10), async {
10707            loop {
10708                if let Ok(WsMessage::Text(text)) = rx.try_recv()
10709                    && text.as_str() == "recovered"
10710                {
10711                    return;
10712                }
10713
10714                tokio::time::sleep(Duration::from_millis(10)).await;
10715            }
10716        })
10717        .await;
10718
10719        let messages = capture.messages();
10720        client.disconnect().await;
10721        server.abort();
10722
10723        recovered.expect("client should recover after the retryable upgrade rejection");
10724
10725        let errors: Vec<_> = messages
10726            .iter()
10727            .filter(|(level, _)| *level == Level::Error)
10728            .collect();
10729        assert!(
10730            errors.is_empty(),
10731            "a retryable rejection must not log at ERROR, was: {errors:?}"
10732        );
10733
10734        // Asserted positively so the ERROR assertion above cannot pass by capturing nothing.
10735        assert!(
10736            messages.iter().any(|(level, message)| *level == Level::Warn
10737                && message.contains(&format!(
10738                    "Sockudo handshake rejected by 127.0.0.1:{port} with retryable status 502"
10739                ))),
10740            "expected a warning naming the retryable rejection, was: {messages:?}"
10741        );
10742
10743        for secret in [
10744            "user-secret",
10745            "password-secret",
10746            "path-secret",
10747            "query-secret",
10748        ] {
10749            assert!(
10750                messages
10751                    .iter()
10752                    .all(|(_, message)| !message.contains(secret)),
10753                "logs must not contain {secret}, was: {messages:?}"
10754            );
10755        }
10756    }
10757
10758    #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
10759    #[rstest]
10760    #[tokio::test]
10761    async fn test_sockudo_backend_replays_leftover_without_custom_headers() {
10762        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
10763        let port = listener.local_addr().unwrap().port();
10764        let server = task::spawn(async move {
10765            if let Ok((mut stream, _)) = listener.accept().await {
10766                let request = read_http_request(&mut stream).await;
10767                let request = String::from_utf8(request).unwrap();
10768                let sec_websocket_key = extract_header(&request, "Sec-WebSocket-Key").unwrap();
10769                let accept = sockudo_handshake::generate_accept_key(sec_websocket_key);
10770                let mut response = format!(
10771                    concat!(
10772                        "HTTP/1.1 101 Switching Protocols\r\n",
10773                        "Upgrade: websocket\r\n",
10774                        "Connection: Upgrade\r\n",
10775                        "Sec-WebSocket-Accept: {}\r\n",
10776                        "\r\n",
10777                    ),
10778                    accept
10779                )
10780                .into_bytes();
10781                response.extend_from_slice(b"\x81\x05hello");
10782                stream.write_all(&response).await.unwrap();
10783            }
10784        });
10785
10786        let (handler, mut rx) = channel_message_handler();
10787
10788        let config = WebSocketConfig {
10789            url: format!("ws://127.0.0.1:{port}/ws"),
10790            headers: vec![],
10791            heartbeat_interval_secs: None,
10792            heartbeat_payload: None,
10793            connect_timeout_ms: Some(2_000),
10794            reconnect_delay_initial_ms: Some(50),
10795            reconnect_delay_max_ms: Some(100),
10796            reconnect_backoff_factor: Some(1.0),
10797            reconnect_jitter_ms: Some(0),
10798            reconnect_max_attempts: None,
10799            heartbeat_timeout_secs: None,
10800            idle_timeout_ms: None,
10801            writer_capacity: None,
10802            backend: TransportBackend::Sockudo,
10803            proxy_url: None,
10804            max_message_size_bytes: None,
10805            max_frame_size_bytes: None,
10806        };
10807
10808        let client = WebSocketClient::builder()
10809            .config(config)
10810            .message_handler(handler)
10811            .connect()
10812            .await
10813            .expect("sockudo connect without custom headers");
10814
10815        let received = tokio::time::timeout(Duration::from_secs(3), async {
10816            loop {
10817                if let Ok(msg) = rx.try_recv() {
10818                    return msg;
10819                }
10820
10821                tokio::time::sleep(Duration::from_millis(10)).await;
10822            }
10823        })
10824        .await
10825        .expect("did not receive leftover frame before timeout");
10826
10827        match received {
10828            WsMessage::Text(t) => assert_eq!(t.as_str(), "hello"),
10829            other => panic!("expected text, was {other:?}"),
10830        }
10831
10832        client.disconnect().await;
10833        tokio::time::timeout(Duration::from_secs(3), server)
10834            .await
10835            .expect("server did not close before timeout")
10836            .unwrap();
10837    }
10838
10839    #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
10840    #[rstest]
10841    #[tokio::test]
10842    async fn test_sockudo_backend_sends_custom_headers() {
10843        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
10844        let port = listener.local_addr().unwrap().port();
10845        let server = task::spawn(async move {
10846            if let Ok((stream, _)) = listener.accept().await {
10847                let callback = HeaderAssertCallback {
10848                    key: "X-Test".to_string(),
10849                    value: HeaderValue::from_static("value"),
10850                };
10851
10852                if let Ok(mut ws) = accept_hdr_async(stream, callback).await {
10853                    while let Some(Ok(msg)) = ws.next().await {
10854                        if msg.is_text() || msg.is_binary() {
10855                            if ws.send(msg).await.is_err() {
10856                                break;
10857                            }
10858
10859                            continue;
10860                        }
10861
10862                        if msg.is_close() {
10863                            let _ = ws.close(None).await;
10864                            break;
10865                        }
10866                    }
10867                }
10868            }
10869        });
10870
10871        let (handler, mut rx) = channel_message_handler();
10872
10873        let config = WebSocketConfig {
10874            url: format!("ws://127.0.0.1:{port}"),
10875            headers: vec![("X-Test".to_string(), "value".to_string())],
10876            heartbeat_interval_secs: None,
10877            heartbeat_payload: None,
10878            connect_timeout_ms: Some(2_000),
10879            reconnect_delay_initial_ms: Some(50),
10880            reconnect_delay_max_ms: Some(100),
10881            reconnect_backoff_factor: Some(1.0),
10882            reconnect_jitter_ms: Some(0),
10883            reconnect_max_attempts: None,
10884            heartbeat_timeout_secs: None,
10885            idle_timeout_ms: None,
10886            writer_capacity: None,
10887            backend: TransportBackend::Sockudo,
10888            proxy_url: None,
10889            max_message_size_bytes: None,
10890            max_frame_size_bytes: None,
10891        };
10892
10893        let client = WebSocketClient::builder()
10894            .config(config)
10895            .message_handler(handler)
10896            .connect()
10897            .await
10898            .expect("sockudo connect with custom headers");
10899
10900        client.send_text("ping".to_string(), None).await.unwrap();
10901
10902        let received = tokio::time::timeout(Duration::from_secs(3), async {
10903            loop {
10904                if let Ok(msg) = rx.try_recv() {
10905                    return msg;
10906                }
10907
10908                tokio::time::sleep(Duration::from_millis(10)).await;
10909            }
10910        })
10911        .await
10912        .expect("did not receive echo before timeout");
10913
10914        match received {
10915            WsMessage::Text(t) => assert_eq!(t.as_str(), "ping"),
10916            other => panic!("expected text, was {other:?}"),
10917        }
10918
10919        client.disconnect().await;
10920        tokio::time::timeout(Duration::from_secs(3), server)
10921            .await
10922            .expect("server did not close before timeout")
10923            .unwrap();
10924    }
10925
10926    #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
10927    #[rstest]
10928    #[tokio::test]
10929    async fn test_sockudo_backend_round_trip_text() {
10930        // tokio-tungstenite test peer paired with a sockudo client.
10931        let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
10932        let port = listener.local_addr().unwrap().port();
10933        let server = task::spawn(async move {
10934            if let Ok((stream, _)) = listener.accept().await
10935                && let Ok(mut ws) = accept_async(stream).await
10936            {
10937                while let Some(Ok(msg)) = ws.next().await {
10938                    match msg {
10939                        WsMessage::Text(_) | WsMessage::Binary(_) => {
10940                            if ws.send(msg).await.is_err() {
10941                                break;
10942                            }
10943                        }
10944                        WsMessage::Close(_) => {
10945                            let _ = ws.close(None).await;
10946                            break;
10947                        }
10948                        _ => {}
10949                    }
10950                }
10951            }
10952        });
10953
10954        let (handler, mut rx) = channel_message_handler();
10955
10956        let config = WebSocketConfig {
10957            url: format!("ws://127.0.0.1:{port}"),
10958            headers: vec![],
10959            heartbeat_interval_secs: None,
10960            heartbeat_payload: None,
10961            connect_timeout_ms: Some(2_000),
10962            reconnect_delay_initial_ms: Some(50),
10963            reconnect_delay_max_ms: Some(100),
10964            reconnect_backoff_factor: Some(1.0),
10965            reconnect_jitter_ms: Some(0),
10966            reconnect_max_attempts: None,
10967            heartbeat_timeout_secs: None,
10968            idle_timeout_ms: None,
10969            writer_capacity: None,
10970            backend: TransportBackend::Sockudo,
10971            proxy_url: None,
10972            max_message_size_bytes: None,
10973            max_frame_size_bytes: None,
10974        };
10975
10976        let client = WebSocketClient::builder()
10977            .config(config)
10978            .message_handler(handler)
10979            .connect()
10980            .await
10981            .expect("sockudo connect");
10982
10983        client.send_text("ping".to_string(), None).await.unwrap();
10984
10985        let received = tokio::time::timeout(Duration::from_secs(3), async {
10986            loop {
10987                if let Ok(msg) = rx.try_recv() {
10988                    return msg;
10989                }
10990
10991                tokio::time::sleep(Duration::from_millis(10)).await;
10992            }
10993        })
10994        .await
10995        .expect("did not receive echo before timeout");
10996
10997        match received {
10998            WsMessage::Text(t) => assert_eq!(t.as_str(), "ping"),
10999            other => panic!("expected text, was {other:?}"),
11000        }
11001
11002        client.disconnect().await;
11003        server.abort();
11004    }
11005
11006    #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
11007    #[rstest]
11008    #[case::ws_default_port("ws://example.com/ws", "example.com", "example.com", 80, "/ws", false)]
11009    #[case::wss_default_port(
11010        "wss://example.com/ws",
11011        "example.com",
11012        "example.com",
11013        443,
11014        "/ws",
11015        true
11016    )]
11017    // url::Url normalizes explicit default ports (`:80` for ws, `:443` for wss)
11018    // away, so `parsed.port()` reports `None` here and Host stays unqualified.
11019    #[case::ws_explicit_default(
11020        "ws://example.com:80/ws",
11021        "example.com",
11022        "example.com",
11023        80,
11024        "/ws",
11025        false
11026    )]
11027    #[case::ws_non_default(
11028        "ws://example.com:8443/feed",
11029        "example.com",
11030        "example.com:8443",
11031        8443,
11032        "/feed",
11033        false
11034    )]
11035    #[case::wss_non_default(
11036        "wss://example.com:9443/feed",
11037        "example.com",
11038        "example.com:9443",
11039        9443,
11040        "/feed",
11041        true
11042    )]
11043    #[case::root_path(
11044        "ws://example.com:9000/",
11045        "example.com",
11046        "example.com:9000",
11047        9000,
11048        "/",
11049        false
11050    )]
11051    #[case::query_string(
11052        "ws://example.com/feed?token=abc&channel=trades",
11053        "example.com",
11054        "example.com",
11055        80,
11056        "/feed?token=abc&channel=trades",
11057        false
11058    )]
11059    // IPv6: bare host strips brackets for DNS/TCP/SNI; Host header keeps them.
11060    #[case::ipv6_default("ws://[::1]/feed", "::1", "[::1]", 80, "/feed", false)]
11061    #[case::ipv6_explicit_port("ws://[::1]:9000/feed", "::1", "[::1]:9000", 9000, "/feed", false)]
11062    #[case::ipv6_wss(
11063        "wss://[2001:db8::1]:8443/",
11064        "2001:db8::1",
11065        "[2001:db8::1]:8443",
11066        8443,
11067        "/",
11068        true
11069    )]
11070    fn sockudo_target_parses_url(
11071        #[case] url: &str,
11072        #[case] host: &str,
11073        #[case] host_header: &str,
11074        #[case] port: u16,
11075        #[case] path: &str,
11076        #[case] is_tls: bool,
11077    ) {
11078        let target = super::SockudoTarget::parse(url).expect("parse should succeed");
11079        assert_eq!(target.host, host);
11080        assert_eq!(target.host_header, host_header);
11081        assert_eq!(target.port, port);
11082        assert_eq!(target.path, path);
11083        assert_eq!(target.is_tls, is_tls);
11084    }
11085
11086    #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
11087    #[rstest]
11088    fn sockudo_target_rejects_unsupported_scheme() {
11089        let err = super::SockudoTarget::parse("http://example.com/feed").expect_err("not a ws URL");
11090        let msg = err.to_string();
11091        assert!(
11092            msg.contains("expected ws:// or wss://"),
11093            "unexpected error: {msg}"
11094        );
11095    }
11096
11097    #[cfg(all(feature = "transport-sockudo", not(feature = "turmoil")))]
11098    #[rstest]
11099    fn sockudo_target_rejects_malformed_url() {
11100        const SECRET: &str = "malformed-websocket-url-secret";
11101        let url = format!("not a url {SECRET}");
11102        let err = super::SockudoTarget::parse(&url).expect_err("malformed URL");
11103        let message = err.to_string();
11104
11105        assert!(
11106            matches!(err, super::TransportError::InvalidUrl(_)),
11107            "expected InvalidUrl, was: {err:?}"
11108        );
11109        assert!(!message.contains(SECRET));
11110        assert!(!message.contains(&url));
11111    }
11112}
11113
11114#[cfg(test)]
11115mod property_tests {
11116    use std::{
11117        collections::{HashSet, VecDeque},
11118        sync::{Arc, OnceLock, atomic::AtomicBool},
11119    };
11120
11121    use proptest::prelude::*;
11122    use rstest::rstest;
11123
11124    use super::{super::auth::AuthResultReceiver, *};
11125
11126    const AUTH_FAILED: &str = "model auth failed";
11127
11128    #[derive(Debug, Clone)]
11129    enum ReconnectBufferTraceOp {
11130        BeginAuth,
11131        AuthSucceeds,
11132        AuthFails,
11133        AuthInvalidates,
11134        ReconnectStarts,
11135        ReconnectCompletes,
11136        BufferedMessage(u8),
11137    }
11138
11139    #[derive(Debug, Clone, Copy, PartialEq, Eq)]
11140    enum ModelConnectionMode {
11141        Active,
11142        Reconnect,
11143    }
11144
11145    #[derive(Debug, Clone, Copy)]
11146    enum ExpectedReconnectBufferAction {
11147        Drain,
11148        Wait,
11149        Discard,
11150    }
11151
11152    #[derive(Debug)]
11153    struct ReconnectBufferModel {
11154        mode: ModelConnectionMode,
11155        auth_state: AuthState,
11156        buffer: VecDeque<String>,
11157        released: Vec<String>,
11158        discarded: Vec<String>,
11159        live_sent: Vec<String>,
11160        handler_controls: Vec<&'static str>,
11161        next_message_index: usize,
11162    }
11163
11164    impl ReconnectBufferModel {
11165        fn new() -> Self {
11166            Self {
11167                mode: ModelConnectionMode::Active,
11168                auth_state: AuthState::Unauthenticated,
11169                buffer: VecDeque::new(),
11170                released: Vec::new(),
11171                discarded: Vec::new(),
11172                live_sent: Vec::new(),
11173                handler_controls: Vec::new(),
11174                next_message_index: 0,
11175            }
11176        }
11177
11178        fn next_payload(&mut self, raw: u8) -> String {
11179            let payload = format!("message-{}-{raw}", self.next_message_index);
11180            self.next_message_index += 1;
11181            payload
11182        }
11183
11184        fn expected_action(&self, waits_for_auth: bool) -> ExpectedReconnectBufferAction {
11185            if !waits_for_auth {
11186                return ExpectedReconnectBufferAction::Drain;
11187            }
11188
11189            match self.auth_state {
11190                AuthState::Authenticated => ExpectedReconnectBufferAction::Drain,
11191                AuthState::Failed => ExpectedReconnectBufferAction::Discard,
11192                AuthState::Unauthenticated => ExpectedReconnectBufferAction::Wait,
11193            }
11194        }
11195    }
11196
11197    fn reconnect_buffer_trace_op_strategy() -> impl Strategy<Value = ReconnectBufferTraceOp> {
11198        prop_oneof![
11199            Just(ReconnectBufferTraceOp::BeginAuth),
11200            Just(ReconnectBufferTraceOp::AuthSucceeds),
11201            Just(ReconnectBufferTraceOp::AuthFails),
11202            Just(ReconnectBufferTraceOp::AuthInvalidates),
11203            Just(ReconnectBufferTraceOp::ReconnectStarts),
11204            Just(ReconnectBufferTraceOp::ReconnectCompletes),
11205            any::<u8>().prop_map(ReconnectBufferTraceOp::BufferedMessage),
11206        ]
11207    }
11208
11209    fn reconnect_buffer_actions_match(
11210        actual: ReconnectBufferAction,
11211        expected: ExpectedReconnectBufferAction,
11212    ) -> bool {
11213        matches!(
11214            (actual, expected),
11215            (
11216                ReconnectBufferAction::Drain,
11217                ExpectedReconnectBufferAction::Drain
11218            ) | (
11219                ReconnectBufferAction::Wait,
11220                ExpectedReconnectBufferAction::Wait
11221            ) | (
11222                ReconnectBufferAction::Discard,
11223                ExpectedReconnectBufferAction::Discard
11224            )
11225        )
11226    }
11227
11228    fn apply_ready_reconnect_buffer_action(
11229        model: &mut ReconnectBufferModel,
11230        reconnect_buffer_waits_for_auth: &AtomicBool,
11231        auth_tracker: &Arc<OnceLock<AuthTracker>>,
11232        waits_for_auth: bool,
11233        step: usize,
11234        op: &ReconnectBufferTraceOp,
11235    ) -> Result<(), TestCaseError> {
11236        if model.mode != ModelConnectionMode::Active || model.buffer.is_empty() {
11237            return Ok(());
11238        }
11239
11240        let expected = model.expected_action(waits_for_auth);
11241        let actual = WebSocketClientInner::can_drain_reconnect_buffer(
11242            reconnect_buffer_waits_for_auth,
11243            auth_tracker,
11244        );
11245
11246        prop_assert!(
11247            reconnect_buffer_actions_match(actual, expected),
11248            "reconnect buffer action mismatch at step {}, op {:?}, waits_for_auth={}, auth_state={:?}",
11249            step,
11250            op,
11251            waits_for_auth,
11252            model.auth_state
11253        );
11254
11255        match expected {
11256            ExpectedReconnectBufferAction::Drain => {
11257                model.released.extend(model.buffer.drain(..));
11258            }
11259            ExpectedReconnectBufferAction::Wait => {}
11260            ExpectedReconnectBufferAction::Discard => {
11261                model.discarded.extend(model.buffer.drain(..));
11262            }
11263        }
11264
11265        Ok(())
11266    }
11267
11268    fn assert_reconnected_control_stays_separate(
11269        model: &ReconnectBufferModel,
11270        step: usize,
11271    ) -> Result<(), TestCaseError> {
11272        prop_assert!(
11273            model
11274                .handler_controls
11275                .iter()
11276                .all(|message| *message == RECONNECTED),
11277            "handler control stream contained a non-RECONNECTED message at step {}",
11278            step
11279        );
11280        prop_assert!(
11281            !model.buffer.iter().any(|message| message == RECONNECTED),
11282            "RECONNECTED control message entered reconnect buffer at step {}",
11283            step
11284        );
11285        prop_assert!(
11286            !model.released.iter().any(|message| message == RECONNECTED),
11287            "RECONNECTED control message entered replayed messages at step {}",
11288            step
11289        );
11290        prop_assert!(
11291            !model.discarded.iter().any(|message| message == RECONNECTED),
11292            "RECONNECTED control message entered discarded messages at step {}",
11293            step
11294        );
11295        prop_assert!(
11296            !model.live_sent.iter().any(|message| message == RECONNECTED),
11297            "RECONNECTED control message entered application sends at step {}",
11298            step
11299        );
11300
11301        Ok(())
11302    }
11303
11304    fn assert_messages_accounted_once(
11305        model: &ReconnectBufferModel,
11306        step: usize,
11307    ) -> Result<(), TestCaseError> {
11308        let mut seen = HashSet::new();
11309
11310        for message in model
11311            .released
11312            .iter()
11313            .chain(model.discarded.iter())
11314            .chain(model.buffer.iter())
11315            .chain(model.live_sent.iter())
11316        {
11317            prop_assert!(
11318                seen.insert(message.as_str()),
11319                "message {} appeared more than once at step {}",
11320                message,
11321                step
11322            );
11323        }
11324
11325        Ok(())
11326    }
11327
11328    fn apply_reconnect_buffer_trace_op(
11329        model: &mut ReconnectBufferModel,
11330        tracker: &AuthTracker,
11331        auth_receivers: &mut Vec<AuthResultReceiver>,
11332        op: &ReconnectBufferTraceOp,
11333    ) -> Result<(), TestCaseError> {
11334        match op {
11335            ReconnectBufferTraceOp::BeginAuth => {
11336                auth_receivers.push(tracker.begin());
11337                model.auth_state = AuthState::Unauthenticated;
11338            }
11339            ReconnectBufferTraceOp::AuthSucceeds => {
11340                tracker.succeed();
11341                model.auth_state = AuthState::Authenticated;
11342            }
11343            ReconnectBufferTraceOp::AuthFails => {
11344                tracker.fail(AUTH_FAILED);
11345                model.auth_state = AuthState::Failed;
11346            }
11347            ReconnectBufferTraceOp::AuthInvalidates => {
11348                tracker.invalidate();
11349
11350                if model.auth_state == AuthState::Authenticated {
11351                    model.auth_state = AuthState::Unauthenticated;
11352                }
11353            }
11354            ReconnectBufferTraceOp::ReconnectStarts => {
11355                tracker.invalidate();
11356
11357                if model.auth_state == AuthState::Authenticated {
11358                    model.auth_state = AuthState::Unauthenticated;
11359                }
11360
11361                model.mode = ModelConnectionMode::Reconnect;
11362            }
11363            ReconnectBufferTraceOp::ReconnectCompletes => {
11364                model.mode = ModelConnectionMode::Active;
11365                model.handler_controls.push(RECONNECTED);
11366            }
11367            ReconnectBufferTraceOp::BufferedMessage(raw) => {
11368                let payload = model.next_payload(*raw);
11369                prop_assert_ne!(payload.as_str(), RECONNECTED);
11370
11371                if model.mode == ModelConnectionMode::Reconnect {
11372                    model.buffer.push_back(payload);
11373                } else {
11374                    model.live_sent.push(payload);
11375                }
11376            }
11377        }
11378
11379        Ok(())
11380    }
11381
11382    proptest! {
11383        #![proptest_config(ProptestConfig::with_cases(256))]
11384
11385        /// Property: reconnect-buffer traces match the auth-gated release and
11386        /// discard model, and `RECONNECTED` remains a separate control signal.
11387        #[rstest]
11388        fn test_reconnect_buffer_trace_matches_auth_gate_model(
11389            waits_for_auth in any::<bool>(),
11390            ops in proptest::collection::vec(reconnect_buffer_trace_op_strategy(), 1..100)
11391        ) {
11392            let auth_tracker = Arc::new(OnceLock::new());
11393            let reconnect_buffer_waits_for_auth = AtomicBool::new(waits_for_auth);
11394            let tracker = AuthTracker::new();
11395            auth_tracker.set(tracker.clone()).unwrap();
11396            let mut auth_receivers = Vec::new();
11397            let mut model = ReconnectBufferModel::new();
11398
11399            for (step, op) in ops.iter().enumerate() {
11400                apply_reconnect_buffer_trace_op(
11401                    &mut model,
11402                    &tracker,
11403                    &mut auth_receivers,
11404                    op,
11405                )?;
11406
11407                prop_assert_eq!(
11408                    tracker.auth_state(),
11409                    model.auth_state,
11410                    "auth state mismatch at step {}, op {:?}",
11411                    step,
11412                    op
11413                );
11414
11415                apply_ready_reconnect_buffer_action(
11416                    &mut model,
11417                    &reconnect_buffer_waits_for_auth,
11418                    &auth_tracker,
11419                    waits_for_auth,
11420                    step,
11421                    op,
11422                )?;
11423                assert_reconnected_control_stays_separate(&model, step)?;
11424                prop_assert_eq!(
11425                    model.handler_controls.len(),
11426                    ops[..=step]
11427                        .iter()
11428                        .filter(|op| matches!(op, ReconnectBufferTraceOp::ReconnectCompletes))
11429                        .count(),
11430                    "handler control count mismatch at step {}",
11431                    step
11432                );
11433                assert_messages_accounted_once(&model, step)?;
11434            }
11435        }
11436
11437        /// Property: successful re-authentication releases buffered messages
11438        /// exactly once when replay is configured to wait for auth.
11439        #[rstest]
11440        fn test_reconnect_buffer_releases_after_auth_success_once(
11441            payloads in proptest::collection::vec(any::<u8>(), 1..32),
11442            extra_success_ticks in 0usize..16
11443        ) {
11444            let auth_tracker = Arc::new(OnceLock::new());
11445            let reconnect_buffer_waits_for_auth = AtomicBool::new(true);
11446            let tracker = AuthTracker::new();
11447            auth_tracker.set(tracker.clone()).unwrap();
11448            let mut auth_receivers = Vec::new();
11449            let mut model = ReconnectBufferModel::new();
11450
11451            apply_reconnect_buffer_trace_op(
11452                &mut model,
11453                &tracker,
11454                &mut auth_receivers,
11455                &ReconnectBufferTraceOp::ReconnectStarts,
11456            )?;
11457            apply_reconnect_buffer_trace_op(
11458                &mut model,
11459                &tracker,
11460                &mut auth_receivers,
11461                &ReconnectBufferTraceOp::BeginAuth,
11462            )?;
11463
11464            for payload in payloads {
11465                apply_reconnect_buffer_trace_op(
11466                    &mut model,
11467                    &tracker,
11468                    &mut auth_receivers,
11469                    &ReconnectBufferTraceOp::BufferedMessage(payload),
11470                )?;
11471            }
11472
11473            let buffered_len = model.buffer.len();
11474            apply_reconnect_buffer_trace_op(
11475                &mut model,
11476                &tracker,
11477                &mut auth_receivers,
11478                &ReconnectBufferTraceOp::ReconnectCompletes,
11479            )?;
11480            apply_ready_reconnect_buffer_action(
11481                &mut model,
11482                &reconnect_buffer_waits_for_auth,
11483                &auth_tracker,
11484                true,
11485                0,
11486                &ReconnectBufferTraceOp::ReconnectCompletes,
11487            )?;
11488
11489            prop_assert_eq!(model.released.len(), 0);
11490            prop_assert_eq!(model.buffer.len(), buffered_len);
11491
11492            apply_reconnect_buffer_trace_op(
11493                &mut model,
11494                &tracker,
11495                &mut auth_receivers,
11496                &ReconnectBufferTraceOp::AuthSucceeds,
11497            )?;
11498            apply_ready_reconnect_buffer_action(
11499                &mut model,
11500                &reconnect_buffer_waits_for_auth,
11501                &auth_tracker,
11502                true,
11503                1,
11504                &ReconnectBufferTraceOp::AuthSucceeds,
11505            )?;
11506
11507            prop_assert_eq!(model.released.len(), buffered_len);
11508            prop_assert!(model.buffer.is_empty());
11509            assert_messages_accounted_once(&model, 1)?;
11510
11511            for tick in 0..extra_success_ticks {
11512                apply_reconnect_buffer_trace_op(
11513                    &mut model,
11514                    &tracker,
11515                    &mut auth_receivers,
11516                    &ReconnectBufferTraceOp::AuthSucceeds,
11517                )?;
11518                apply_ready_reconnect_buffer_action(
11519                    &mut model,
11520                    &reconnect_buffer_waits_for_auth,
11521                    &auth_tracker,
11522                    true,
11523                    tick + 2,
11524                    &ReconnectBufferTraceOp::AuthSucceeds,
11525                )?;
11526                prop_assert_eq!(
11527                    model.released.len(),
11528                    buffered_len,
11529                    "buffered messages replayed more than once at tick {}",
11530                    tick
11531                );
11532            }
11533        }
11534
11535        /// Property: auth failure discards messages buffered before or after
11536        /// that failure, and later auth success does not replay discarded data.
11537        #[rstest]
11538        fn test_reconnect_buffer_discards_after_auth_failure(
11539            before_failure_payloads in proptest::collection::vec(any::<u8>(), 0..16),
11540            after_failure_payloads in proptest::collection::vec(any::<u8>(), 1..16),
11541            later_success_ticks in 0usize..16
11542        ) {
11543            let auth_tracker = Arc::new(OnceLock::new());
11544            let reconnect_buffer_waits_for_auth = AtomicBool::new(true);
11545            let tracker = AuthTracker::new();
11546            auth_tracker.set(tracker.clone()).unwrap();
11547            let mut auth_receivers = Vec::new();
11548            let mut model = ReconnectBufferModel::new();
11549
11550            apply_reconnect_buffer_trace_op(
11551                &mut model,
11552                &tracker,
11553                &mut auth_receivers,
11554                &ReconnectBufferTraceOp::ReconnectStarts,
11555            )?;
11556            apply_reconnect_buffer_trace_op(
11557                &mut model,
11558                &tracker,
11559                &mut auth_receivers,
11560                &ReconnectBufferTraceOp::BeginAuth,
11561            )?;
11562
11563            for payload in before_failure_payloads {
11564                apply_reconnect_buffer_trace_op(
11565                    &mut model,
11566                    &tracker,
11567                    &mut auth_receivers,
11568                    &ReconnectBufferTraceOp::BufferedMessage(payload),
11569                )?;
11570            }
11571
11572            apply_reconnect_buffer_trace_op(
11573                &mut model,
11574                &tracker,
11575                &mut auth_receivers,
11576                &ReconnectBufferTraceOp::AuthFails,
11577            )?;
11578
11579            for payload in after_failure_payloads {
11580                apply_reconnect_buffer_trace_op(
11581                    &mut model,
11582                    &tracker,
11583                    &mut auth_receivers,
11584                    &ReconnectBufferTraceOp::BufferedMessage(payload),
11585                )?;
11586            }
11587
11588            let buffered_len = model.buffer.len();
11589            apply_reconnect_buffer_trace_op(
11590                &mut model,
11591                &tracker,
11592                &mut auth_receivers,
11593                &ReconnectBufferTraceOp::ReconnectCompletes,
11594            )?;
11595            apply_ready_reconnect_buffer_action(
11596                &mut model,
11597                &reconnect_buffer_waits_for_auth,
11598                &auth_tracker,
11599                true,
11600                0,
11601                &ReconnectBufferTraceOp::ReconnectCompletes,
11602            )?;
11603
11604            prop_assert_eq!(model.discarded.len(), buffered_len);
11605            prop_assert!(model.released.is_empty());
11606            prop_assert!(model.buffer.is_empty());
11607            assert_messages_accounted_once(&model, 0)?;
11608
11609            for tick in 0..later_success_ticks {
11610                apply_reconnect_buffer_trace_op(
11611                    &mut model,
11612                    &tracker,
11613                    &mut auth_receivers,
11614                    &ReconnectBufferTraceOp::BeginAuth,
11615                )?;
11616                apply_reconnect_buffer_trace_op(
11617                    &mut model,
11618                    &tracker,
11619                    &mut auth_receivers,
11620                    &ReconnectBufferTraceOp::AuthSucceeds,
11621                )?;
11622                apply_ready_reconnect_buffer_action(
11623                    &mut model,
11624                    &reconnect_buffer_waits_for_auth,
11625                    &auth_tracker,
11626                    true,
11627                    tick + 1,
11628                    &ReconnectBufferTraceOp::AuthSucceeds,
11629                )?;
11630                prop_assert!(
11631                    model.released.is_empty(),
11632                    "discarded messages replayed after later auth success at tick {}",
11633                    tick
11634                );
11635            }
11636        }
11637    }
11638}
11639
11640#[cfg(test)]
11641#[cfg(feature = "turmoil")]
11642mod turmoil_tests {
11643    use std::{sync::Arc, time::Duration};
11644
11645    use futures_util::{SinkExt, StreamExt};
11646    use nautilus_common::testing::wait_until_async;
11647    use rstest::rstest;
11648    use tokio_tungstenite::{accept_async, tungstenite::Message as WsMessage};
11649    use turmoil::{Builder, net};
11650
11651    use super::*;
11652    use crate::websocket::types::channel_message_handler;
11653
11654    const AUTH_BUFFER_WAIT_SEED: u64 = 0xA17B_0001;
11655    const AUTH_BUFFER_DISCARD_SEED: u64 = 0xA17B_0002;
11656
11657    fn seeded_turmoil_builder(seed: u64) -> Builder {
11658        let mut builder = Builder::new();
11659        builder.rng_seed(seed);
11660        builder
11661    }
11662
11663    #[rstest]
11664    fn test_turmoil_reconnect_buffer_waits_for_auth() {
11665        let mut sim = seeded_turmoil_builder(AUTH_BUFFER_WAIT_SEED).build();
11666        let messages = Arc::new(tokio::sync::Mutex::new(Vec::new()));
11667        let server_messages = Arc::clone(&messages);
11668        sim.host("server", move || {
11669            let messages = Arc::clone(&server_messages);
11670            auth_buffer_server(messages)
11671        });
11672
11673        sim.client("client", async move {
11674            let tracker = AuthTracker::new();
11675            let (handler, _rx) = channel_message_handler();
11676            let client = WebSocketClient::builder()
11677                .config(turmoil_websocket_config())
11678                .message_handler(handler)
11679                .connect()
11680                .await
11681                .expect("Should connect");
11682
11683            client.set_auth_tracker(tracker.clone(), true);
11684            assert!(client.is_active(), "Client should start active");
11685
11686            wait_until_async(
11687                || async { client.is_reconnecting() },
11688                Duration::from_secs(3),
11689            )
11690            .await;
11691
11692            client
11693                .writer_tx
11694                .send(WriterCommand::Send(Message::Text("stale".into())))
11695                .unwrap();
11696
11697            wait_until_async(|| async { client.is_active() }, Duration::from_secs(3)).await;
11698
11699            let _auth_receiver = tracker.begin();
11700
11701            tokio::time::sleep(Duration::from_millis(300)).await;
11702            assert!(
11703                messages.lock().await.is_empty(),
11704                "buffered messages should wait for auth after reconnect"
11705            );
11706
11707            tracker.succeed();
11708
11709            wait_until_async(
11710                || {
11711                    let messages = Arc::clone(&messages);
11712                    async move { messages.lock().await.as_slice() == ["stale"] }
11713                },
11714                Duration::from_secs(3),
11715            )
11716            .await;
11717
11718            assert_eq!(messages.lock().await.as_slice(), ["stale"]);
11719
11720            client.disconnect().await;
11721            assert!(client.is_disconnected());
11722
11723            Ok(())
11724        });
11725
11726        sim.run().unwrap();
11727    }
11728
11729    #[rstest]
11730    fn test_turmoil_reconnect_buffer_discards_after_auth_failure() {
11731        let mut sim = seeded_turmoil_builder(AUTH_BUFFER_DISCARD_SEED).build();
11732        let messages = Arc::new(tokio::sync::Mutex::new(Vec::new()));
11733        let server_messages = Arc::clone(&messages);
11734        sim.host("server", move || {
11735            let messages = Arc::clone(&server_messages);
11736            auth_buffer_server(messages)
11737        });
11738
11739        sim.client("client", async move {
11740            let tracker = AuthTracker::new();
11741            let (handler, _rx) = channel_message_handler();
11742            let client = WebSocketClient::builder()
11743                .config(turmoil_websocket_config())
11744                .message_handler(handler)
11745                .connect()
11746                .await
11747                .expect("Should connect");
11748
11749            client.set_auth_tracker(tracker.clone(), true);
11750            assert!(client.is_active(), "Client should start active");
11751
11752            wait_until_async(
11753                || async { client.is_reconnecting() },
11754                Duration::from_secs(3),
11755            )
11756            .await;
11757
11758            client
11759                .writer_tx
11760                .send(WriterCommand::Send(Message::Text("stale".into())))
11761                .unwrap();
11762
11763            wait_until_async(|| async { client.is_active() }, Duration::from_secs(3)).await;
11764
11765            let _auth_receiver = tracker.begin();
11766            tracker.fail("rejected");
11767
11768            tokio::time::sleep(Duration::from_millis(300)).await;
11769            assert!(
11770                messages.lock().await.is_empty(),
11771                "buffered messages should be discarded after auth failure"
11772            );
11773
11774            let _retry_auth_receiver = tracker.begin();
11775            tracker.succeed();
11776
11777            tokio::time::sleep(Duration::from_millis(300)).await;
11778            assert!(
11779                messages.lock().await.is_empty(),
11780                "discarded messages should not replay on a later auth success"
11781            );
11782
11783            client.disconnect().await;
11784            assert!(client.is_disconnected());
11785
11786            Ok(())
11787        });
11788
11789        sim.run().unwrap();
11790    }
11791
11792    fn turmoil_websocket_config() -> WebSocketConfig {
11793        WebSocketConfig {
11794            url: "ws://server:8080".to_string(),
11795            headers: vec![],
11796            heartbeat_interval_secs: None,
11797            heartbeat_payload: None,
11798            connect_timeout_ms: Some(5_000),
11799            reconnect_delay_initial_ms: Some(50),
11800            reconnect_delay_max_ms: Some(200),
11801            reconnect_backoff_factor: Some(1.0),
11802            reconnect_jitter_ms: Some(0),
11803            reconnect_max_attempts: None,
11804            heartbeat_timeout_secs: None,
11805            idle_timeout_ms: None,
11806            writer_capacity: None,
11807            backend: TransportBackend::Tungstenite,
11808            proxy_url: None,
11809            max_message_size_bytes: None,
11810            max_frame_size_bytes: None,
11811        }
11812    }
11813
11814    async fn auth_buffer_server(
11815        messages: Arc<tokio::sync::Mutex<Vec<String>>>,
11816    ) -> Result<(), Box<dyn std::error::Error>> {
11817        let listener = net::TcpListener::bind("0.0.0.0:8080").await?;
11818
11819        let (stream, _) = listener.accept().await?;
11820        let mut websocket = accept_async(stream).await?;
11821        let _ = websocket.send(WsMessage::Text("first".into())).await;
11822        drop(websocket);
11823
11824        tokio::time::sleep(Duration::from_millis(200)).await;
11825
11826        let (stream, _) = listener.accept().await?;
11827        let mut websocket = accept_async(stream).await?;
11828
11829        while let Some(msg) = websocket.next().await {
11830            match msg {
11831                Ok(WsMessage::Text(text)) => {
11832                    messages.lock().await.push(text.to_string());
11833                }
11834                Ok(WsMessage::Close(_)) => {
11835                    let _ = websocket.close(None).await;
11836                    break;
11837                }
11838                Ok(_) => {}
11839                Err(_) => break,
11840            }
11841        }
11842
11843        Ok(())
11844    }
11845}