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