nautilus_lighter/websocket/

client.rs

// -------------------------------------------------------------------------------------------------
//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
//  https://nautechsystems.io
//
//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
//  You may not use this file except in compliance with the License.
//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
//
//  Unless required by applicable law or agreed to in writing, software
//  distributed under the License is distributed on an "AS IS" BASIS,
//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
//  See the License for the specific language governing permissions and
//  limitations under the License.
// -------------------------------------------------------------------------------------------------

//! Outer WebSocket client orchestrating connection lifecycle and subscriptions.

use std::{
    fmt::Debug,
    num::NonZeroU32,
    sync::{
        Arc,
        atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
    },
    time::Duration,
};

use arc_swap::ArcSwap;
use dashmap::{DashMap, mapref::entry::Entry};
#[cfg(test)]
use nautilus_common::live::get_runtime;
use nautilus_core::string::secret::{SecretString, redact_option};
#[cfg(test)]
use nautilus_live::book::DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS;
use nautilus_live::{
    SocketControl,
    task::{SharedTaskSlot, TaskJoinOutcome, TaskSlot, finish_task},
};
use nautilus_model::{
    identifiers::{AccountId, InstrumentId},
    instruments::InstrumentAny,
};
use nautilus_network::{
    ConnectionUrl, SocketStateSink,
    http::create_standard_nautilus_headers,
    mode::ConnectionMode,
    websocket::{
        InitialConnectRetryPolicy, SubscriptionState, TransportBackend, WebSocketClient,
        WebSocketConfig, channel_epoch_message_handler,
    },
};
use parking_lot::Mutex;
use tokio_util::sync::CancellationToken;

use crate::{
    common::{
        consts::{
            DISCONNECT_TIMEOUT, HEARTBEAT_INTERVAL, HEARTBEAT_TIMEOUT, RECONNECT_BASE_BACKOFF,
            RECONNECT_MAX_BACKOFF,
        },
        enums::{LighterCandleResolution, LighterEnvironment},
        rate_limit::ws_message_rate_limiter,
        symbol::MarketRegistry,
        urls::lighter_ws_url,
    },
    websocket::{
        error::LighterWsError,
        handler::{FeedHandler, HandlerCommand},
        messages::{LighterMarketSelection, LighterWsChannel, NautilusWsMessage},
    },
};

const RECONNECT_JITTER_MS: u64 = 200;
const RECONNECT_BACKOFF_FACTOR: f64 = 2.0;

#[derive(Clone)]
struct SubscriptionArgs {
    channel: LighterWsChannel,
    auth: Option<SecretString>,
    generation: u64,
}

/// Outer Lighter WebSocket client.
///
/// Orchestrates the connection lifecycle and subscription bookkeeping for the
/// Lighter streaming API. The inner feed handler runs on a dedicated tokio
/// task and exclusively owns the underlying [`WebSocketClient`]; this outer
/// type communicates with it through a command channel and consumes events
/// over an unbounded mpsc.
///
/// Authenticated channels store their auth token in `subscription_args` and
/// replay it verbatim on reconnect. That stored token stays valid because the
/// execution client rotates it on a 6h cadence (inside the venue's 7h TTL) by
/// re-issuing `subscribe_account`, so a reconnect never replays a connect-time,
/// potentially-expired token.
pub struct LighterWebSocketClient {
    url: ConnectionUrl,
    connection_mode: Arc<ArcSwap<AtomicU8>>,
    connection_epoch: Arc<ArcSwap<AtomicU64>>,
    connection_lock: Arc<tokio::sync::Mutex<()>>,
    connection_generation: Arc<AtomicU64>,
    initial_connect_cancellation: Arc<ArcSwap<CancellationToken>>,
    signal: Arc<AtomicBool>,
    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
    out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<NautilusWsMessage>>,
    subscriptions: SubscriptionState,
    subscription_args: Arc<DashMap<String, SubscriptionArgs>>,
    next_subscription_generation: Arc<AtomicU64>,
    instruments: Arc<DashMap<i64, InstrumentAny>>,
    cache_lock: Arc<Mutex<()>>,
    registry: Arc<MarketRegistry>,
    task_handle: TaskSlot<()>,
    transport_backend: TransportBackend,
    ws_timeout_secs: u64,
    book_snapshot_timeout: Duration,
    proxy_url: Option<SecretString>,
    socket_sink: Option<SocketStateSink>,
    socket_control: Option<SocketControl>,
}

#[derive(Debug)]
pub(crate) struct RetainedTaskSlot(SharedTaskSlot<()>);

impl RetainedTaskSlot {
    pub(crate) fn new() -> Self {
        Self(SharedTaskSlot::new())
    }

    pub(crate) fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    pub(crate) async fn finish(&self) -> Result<(), LighterWsError> {
        let Some(outcome) = self.0.finish(DISCONNECT_TIMEOUT, DISCONNECT_TIMEOUT).await else {
            return Ok(());
        };

        match outcome {
            TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => Ok(()),
            TaskJoinOutcome::Failed(error) => Err(LighterWsError::Client(format!(
                "retained WebSocket handler task failed: {error}"
            ))),
            TaskJoinOutcome::Incomplete => Err(LighterWsError::Client(
                "retained WebSocket handler task did not stop after abort".to_string(),
            )),
        }
    }
}

pub(crate) struct TaskRetentionGuard {
    client: Option<LighterWebSocketClient>,
    retained: Arc<RetainedTaskSlot>,
}

impl TaskRetentionGuard {
    pub(crate) fn new(client: LighterWebSocketClient, retained: Arc<RetainedTaskSlot>) -> Self {
        Self {
            client: Some(client),
            retained,
        }
    }

    pub(crate) fn client_mut(&mut self) -> &mut LighterWebSocketClient {
        self.client.as_mut().expect("retention guard is armed")
    }

    pub(crate) fn disarm(mut self) -> LighterWebSocketClient {
        self.client.take().expect("retention guard is armed")
    }
}

impl Drop for TaskRetentionGuard {
    fn drop(&mut self) {
        let Some(client) = self.client.as_mut() else {
            return;
        };
        client.begin_shutdown();
        let slot = client.take_task_slot();
        if slot.is_some() && self.retained.0.try_insert_slot(slot).is_err() {
            log::error!("Lighter retained WebSocket task slot was already occupied");
        }
    }
}

impl Debug for LighterWebSocketClient {
    /// Custom `Debug` that redacts the auth token in `subscription_args`.
    ///
    /// Authenticated channel subscriptions store the venue bearer token
    /// alongside the channel for reconnect replay; deriving `Debug` would
    /// otherwise print the token verbatim.
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        let subscription_topics: Vec<String> = self
            .subscription_args
            .iter()
            .map(|entry| {
                let args = entry.value();
                format!(
                    "topic={} channel={:?} authed={} generation={}",
                    entry.key(),
                    args.channel,
                    args.auth.is_some(),
                    args.generation,
                )
            })
            .collect();

        f.debug_struct(stringify!(LighterWebSocketClient))
            .field("url", &self.url)
            .field("is_active", &self.is_active())
            .field("subscription_count", &self.subscriptions.len())
            .field("subscription_args", &subscription_topics)
            .field("instruments_len", &self.instruments.len())
            .field("transport_backend", &self.transport_backend)
            .field("ws_timeout_secs", &self.ws_timeout_secs)
            .field("book_snapshot_timeout", &self.book_snapshot_timeout)
            .field("proxy_url", &redact_option(self.proxy_url.as_ref()))
            .finish_non_exhaustive()
    }
}

impl Clone for LighterWebSocketClient {
    fn clone(&self) -> Self {
        Self {
            url: self.url.clone(),
            connection_mode: Arc::clone(&self.connection_mode),
            connection_epoch: Arc::clone(&self.connection_epoch),
            connection_lock: Arc::clone(&self.connection_lock),
            connection_generation: Arc::clone(&self.connection_generation),
            initial_connect_cancellation: Arc::clone(&self.initial_connect_cancellation),
            signal: Arc::clone(&self.signal),
            cmd_tx: Arc::clone(&self.cmd_tx),
            out_rx: None,
            subscriptions: self.subscriptions.clone(),
            subscription_args: Arc::clone(&self.subscription_args),
            next_subscription_generation: Arc::clone(&self.next_subscription_generation),
            instruments: Arc::clone(&self.instruments),
            cache_lock: Arc::clone(&self.cache_lock),
            registry: Arc::clone(&self.registry),
            task_handle: TaskSlot::new(),
            transport_backend: self.transport_backend,
            ws_timeout_secs: self.ws_timeout_secs,
            book_snapshot_timeout: self.book_snapshot_timeout,
            proxy_url: self.proxy_url.clone(),
            socket_sink: self.socket_sink.clone(),
            socket_control: self.socket_control.clone(),
        }
    }
}

impl LighterWebSocketClient {
    fn initial_connect_retry_policy() -> InitialConnectRetryPolicy {
        InitialConnectRetryPolicy {
            max_attempts: NonZeroU32::new(5).expect("initial connect attempts must be non-zero"),
            delay_initial: Duration::from_millis(500),
            delay_max: Duration::from_secs(5),
            backoff_factor: 2.0,
            jitter_ms: 250,
        }
    }

    /// Creates a new client without connecting.
    ///
    /// `url` overrides the resolved environment URL when supplied.
    #[must_use]
    pub fn new(
        url: Option<String>,
        environment: LighterEnvironment,
        registry: Arc<MarketRegistry>,
        transport_backend: TransportBackend,
        ws_timeout_secs: u64,
        book_snapshot_timeout: Duration,
        proxy_url: Option<String>,
    ) -> Self {
        let url = url.unwrap_or_else(|| lighter_ws_url(environment).to_string());
        let connection_mode = Arc::new(ArcSwap::new(Arc::new(AtomicU8::new(
            ConnectionMode::Closed as u8,
        ))));
        let connection_epoch = Arc::new(ArcSwap::new(Arc::new(AtomicU64::new(0))));

        let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();

        Self {
            url: url.into(),
            connection_mode,
            connection_epoch,
            connection_lock: Arc::new(tokio::sync::Mutex::new(())),
            connection_generation: Arc::new(AtomicU64::new(0)),
            initial_connect_cancellation: Arc::new(ArcSwap::from_pointee(CancellationToken::new())),
            signal: Arc::new(AtomicBool::new(false)),
            cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
            out_rx: None,
            subscriptions: SubscriptionState::new(':'),
            subscription_args: Arc::new(DashMap::new()),
            next_subscription_generation: Arc::new(AtomicU64::new(1)),
            instruments: Arc::new(DashMap::new()),
            cache_lock: Arc::new(Mutex::new(())),
            registry,
            task_handle: TaskSlot::new(),
            transport_backend,
            ws_timeout_secs,
            book_snapshot_timeout,
            proxy_url: proxy_url.map(SecretString::from),
            socket_sink: None,
            socket_control: None,
        }
    }

    /// Configures socket state reporting for the underlying transport.
    #[must_use]
    pub fn with_state_sink(mut self, state_sink: SocketStateSink) -> Self {
        self.socket_sink = Some(state_sink);
        self
    }

    /// Configures socket state reporting and reconnect control.
    #[must_use]
    pub fn with_socket_control(mut self, control: SocketControl) -> Self {
        self.socket_control = Some(control);
        self
    }

    /// Returns the resolved WebSocket URL.
    #[must_use]
    pub fn url(&self) -> &str {
        self.url.expose()
    }

    /// Returns `true` when the underlying connection is active.
    #[must_use]
    pub fn is_active(&self) -> bool {
        self.connection_mode.load().load(Ordering::Relaxed) == ConnectionMode::Active as u8
    }

    #[must_use]
    pub(crate) fn connection_epoch(&self) -> u64 {
        self.connection_epoch.load().load(Ordering::Acquire)
    }

    pub(crate) fn connection_epoch_atomic(&self) -> Arc<AtomicU64> {
        self.connection_epoch.load_full()
    }

    /// Waits until the underlying connection reports active, or returns an
    /// error after the configured WebSocket timeout.
    ///
    /// Polls [`Self::is_active`] every 10ms. Mirrors the documented
    /// `wait_until_active` contract for adapter WebSocket clients in
    /// `docs/developer_guide/adapters.md`.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection does not reach the active state within the configured
    /// timeout.
    pub async fn wait_until_active(&self) -> Result<(), LighterWsError> {
        let timeout_secs = self.ws_timeout_secs;
        let timeout = Duration::from_secs(timeout_secs);

        tokio::time::timeout(timeout, async {
            while !self.is_active() {
                tokio::time::sleep(Duration::from_millis(10)).await;
            }
        })
        .await
        .map_err(|_| {
            LighterWsError::Client(format!(
                "WebSocket connection timeout after {timeout_secs} seconds"
            ))
        })
    }

    /// Returns the count of confirmed subscriptions.
    #[must_use]
    pub fn subscription_count(&self) -> usize {
        self.subscriptions.len()
    }

    /// Returns a clone of the shared instrument cache.
    #[must_use]
    pub fn instruments_cache(&self) -> Arc<DashMap<i64, InstrumentAny>> {
        Arc::clone(&self.instruments)
    }

    /// Caches a batch of instruments along with their venue `market_index`,
    /// replaying them to the handler if a connection is already established.
    pub fn cache_instruments(&self, instruments: Vec<(i64, InstrumentAny)>) {
        let _guard = self.cache_lock.lock();
        self.instruments.clear();
        for (market_index, instrument) in &instruments {
            self.instruments.insert(*market_index, instrument.clone());
        }
        log::debug!(
            "Lighter instrument cache initialized with {} instruments",
            instruments.len()
        );

        if let Ok(cmd_tx) = self.cmd_tx.try_read() {
            let _ = cmd_tx.send(HandlerCommand::InitializeInstruments(instruments));
        }
    }

    /// Caches a single instrument and pushes it to the handler if connected.
    pub fn cache_instrument(&self, market_index: i64, instrument: InstrumentAny) {
        let _guard = self.cache_lock.lock();
        self.instruments.insert(market_index, instrument.clone());

        if let Ok(cmd_tx) = self.cmd_tx.try_read() {
            let _ = cmd_tx.send(HandlerCommand::UpdateInstrument {
                market_index,
                instrument,
            });
        }
    }

    /// Establishes the WebSocket connection and spawns the feed-handler task.
    /// Classified transient failures retry within the configured WebSocket timeout.
    ///
    /// A fresh connection clears held subscriptions; callers subscribe again for the new session.
    /// Transport reconnects, including `SocketControl` requests, preserve and replay them.
    ///
    /// # Errors
    ///
    /// Returns an error if the connection fails permanently, exhausts its timeout, is cancelled,
    /// or the handler cannot be initialized.
    pub async fn connect(&mut self) -> anyhow::Result<()> {
        self.connect_with_cancellation(CancellationToken::new())
            .await
    }

    pub(crate) async fn connect_with_cancellation(
        &mut self,
        cancellation_token: CancellationToken,
    ) -> anyhow::Result<()> {
        let generation = self.connection_generation.load(Ordering::Acquire);
        let _guard = self.connection_lock.lock().await;

        anyhow::ensure!(
            generation == self.connection_generation.load(Ordering::Acquire),
            "Lighter WebSocket initial connection cancelled",
        );

        if self.is_active() {
            log::warn!("Lighter WebSocket already connected");
            return Ok(());
        }

        if let Some(outcome) = finish_task(
            &mut self.task_handle,
            DISCONNECT_TIMEOUT,
            DISCONNECT_TIMEOUT,
        )
        .await
        {
            match outcome {
                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
                TaskJoinOutcome::Failed(error) => {
                    anyhow::bail!("Lighter WebSocket handler failed: {error}");
                }
                TaskJoinOutcome::Incomplete => {
                    anyhow::bail!("Lighter WebSocket handler did not stop after abort");
                }
            }
        }

        self.signal.store(false, Ordering::Release);
        self.initial_connect_cancellation
            .store(Arc::new(cancellation_token.clone()));

        let (message_handler, raw_rx) = channel_epoch_message_handler();
        let headers = create_standard_nautilus_headers();

        let cfg = WebSocketConfig {
            url: self.url.expose().to_owned(),
            headers,
            heartbeat_interval_secs: Some(HEARTBEAT_INTERVAL.as_secs()),
            heartbeat_payload: None,
            connect_timeout_ms: Some(self.ws_timeout_secs.saturating_mul(1_000).max(1)),
            reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
            reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
            reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
            reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
            reconnect_max_attempts: None,
            heartbeat_timeout_secs: Some(HEARTBEAT_TIMEOUT.as_secs()),
            idle_timeout_ms: None,
            writer_capacity: None,
            backend: self.transport_backend,
            proxy_url: self
                .proxy_url
                .as_ref()
                .map(|value| value.expose_secret().to_owned()),
            max_message_size_bytes: None,
            max_frame_size_bytes: None,
        };

        let connect = WebSocketClient::epoch_builder()
            .config(cfg)
            .epoch_handler(message_handler)
            .rate_limiter(ws_message_rate_limiter(self.url.expose()))
            .initial_connect_retry_policy(Self::initial_connect_retry_policy())
            .cancellation_token(cancellation_token.clone())
            .maybe_state_sink(
                self.socket_control
                    .as_ref()
                    .map(SocketControl::sink)
                    .or_else(|| self.socket_sink.clone()),
            )
            .connect();
        let client =
            match tokio::time::timeout(Duration::from_secs(self.ws_timeout_secs), connect).await {
                Ok(result) => result?,
                Err(_) => anyhow::bail!(
                    "Lighter WebSocket initial connection timeout after {} seconds",
                    self.ws_timeout_secs,
                ),
            };

        if cancellation_token.is_cancelled()
            || generation != self.connection_generation.load(Ordering::Acquire)
        {
            client.disconnect().await;

            anyhow::bail!("Lighter WebSocket initial connection cancelled");
        }

        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<NautilusWsMessage>();

        // Capture the connection-mode atomic before moving `client` into the
        // SetClient command below.
        let connection_mode_atomic = client.connection_mode_atomic();
        let connection_epoch_atomic = client.connection_epoch_atomic();

        // Queue SetClient (and the instrument cache replay) onto the new
        // command channel BEFORE publishing it to clones or marking the
        // connection active. Otherwise a clone observing `is_active()` could
        // race in and send a Subscribe before SetClient lands, and the
        // handler would drop the subscription because `inner == None`.
        let reconnect_handle = client.reconnect_handle();
        if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
            anyhow::bail!("Failed to send SetClient command: {e}");
        }

        if let Some(control) = &self.socket_control {
            control.register(move || reconnect_handle.request_reconnect());
        }

        self.out_rx = Some(out_rx);
        {
            let mut sender = self.cmd_tx.write().await;
            // Only connect takes the sender write lock, so skipped cache sends join this snapshot
            let cache_guard = self.cache_lock.lock();
            let initial_instruments: Vec<(i64, InstrumentAny)> = self
                .instruments
                .iter()
                .map(|entry| (*entry.key(), entry.value().clone()))
                .collect();

            if !initial_instruments.is_empty()
                && let Err(e) =
                    cmd_tx.send(HandlerCommand::InitializeInstruments(initial_instruments))
            {
                log::error!("Failed to send InitializeInstruments: {e}");
            }

            self.subscriptions.clear();
            self.subscription_args.clear();
            *sender = cmd_tx.clone();
            self.connection_epoch.store(connection_epoch_atomic);
            self.connection_mode.store(connection_mode_atomic);
            // Release the sender first so cache writes cannot miss both the snapshot and dispatch
            drop(sender);
            drop(cache_guard);
        }

        log::debug!("Lighter WebSocket connected: {}", self.url.origin());

        let signal = Arc::clone(&self.signal);
        let subscriptions = self.subscriptions.clone();
        let subscription_args = Arc::clone(&self.subscription_args);
        let cmd_tx_for_reconnect = cmd_tx.clone();
        let settlement_currency = self.registry.settlement_currency();
        let book_snapshot_timeout = self.book_snapshot_timeout;

        if let Err(e) = self.task_handle.spawn(async move {
            let mut handler = FeedHandler::new_with_settlement_currency(
                Arc::clone(&signal),
                cmd_rx,
                raw_rx,
                out_tx,
                subscriptions,
                settlement_currency,
                book_snapshot_timeout,
            );

            handler.set_command_sender(cmd_tx_for_reconnect.clone());

            let restore_subscriptions = || {
                if subscription_args.is_empty() {
                    log::debug!("No active Lighter subscriptions to restore after reconnect");
                    return;
                }
                log::debug!(
                    "Restoring {} Lighter subscriptions after reconnect",
                    subscription_args.len(),
                );

                // Replay first; the execution client replaces account tokens after reconnect
                for entry in subscription_args.iter() {
                    let args = entry.value().clone();
                    if let Err(e) = cmd_tx_for_reconnect.send(HandlerCommand::Subscribe {
                        channel: args.channel,
                        auth: args.auth,
                        response_tx: None,
                    }) {
                        log::error!("Failed to resend Lighter subscribe command: {e}");
                    }
                }
            };

            loop {
                match handler.next().await {
                    Some(NautilusWsMessage::Reconnected { connection_epoch }) => {
                        log::debug!("Lighter WebSocket reconnected");
                        restore_subscriptions();

                        if handler
                            .send(NautilusWsMessage::Reconnected { connection_epoch })
                            .is_err()
                        {
                            if handler.is_stopped() {
                                log::debug!("Failed to forward Reconnected (receiver dropped)");
                            } else {
                                log::error!("Failed to forward Reconnected (receiver dropped)");
                            }
                            break;
                        }
                    }
                    Some(msg) => {
                        if handler.send(msg).is_err() {
                            if handler.is_stopped() {
                                log::debug!("Failed to send Lighter message (receiver dropped)");
                            } else {
                                log::error!("Failed to send Lighter message (receiver dropped)");
                            }
                            break;
                        }
                    }
                    None => {
                        if handler.is_stopped() {
                            log::debug!("Lighter handler stop signal observed, exiting loop");
                            break;
                        }
                        log::warn!("Lighter WebSocket stream ended unexpectedly");
                        break;
                    }
                }
            }
            log::debug!("Lighter handler task completed");
        }) {
            self.out_rx = None;
            anyhow::bail!("Failed to start Lighter WebSocket handler task: {e}");
        }
        Ok(())
    }

    /// Disconnects gracefully: signals shutdown, drains the handler, then
    /// awaits the task handle with a timeout.
    ///
    /// # Errors
    ///
    /// Returns an error if the handler task fails or does not stop after being aborted. Shutdown
    /// still clears the connection mode and deregisters the socket control.
    pub async fn disconnect(&mut self) -> Result<(), LighterWsError> {
        self.connection_generation.fetch_add(1, Ordering::AcqRel);
        self.initial_connect_cancellation.load().cancel();

        let _guard = self.connection_lock.lock().await;
        self.initial_connect_cancellation.load().cancel();

        log::debug!("Disconnecting Lighter WebSocket");

        if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
            log::debug!("Failed to send Lighter disconnect command: {e}");
        }
        self.signal.store(true, Ordering::Release);

        let task_result = match finish_task(
            &mut self.task_handle,
            DISCONNECT_TIMEOUT,
            DISCONNECT_TIMEOUT,
        )
        .await
        {
            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => Ok(()),
            Some(TaskJoinOutcome::Failed(error)) => Err(LighterWsError::Client(format!(
                "WebSocket handler task failed: {error}"
            ))),
            Some(TaskJoinOutcome::Incomplete) => Err(LighterWsError::Client(
                "WebSocket handler task did not stop after abort".to_string(),
            )),
        };

        self.connection_mode
            .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));

        if let Some(control) = &self.socket_control {
            control.deregister();
        }
        task_result
    }

    pub(crate) async fn disconnect_with_task_retention(
        self,
        retained: Arc<RetainedTaskSlot>,
    ) -> Result<(), LighterWsError> {
        let mut guard = TaskRetentionGuard::new(self, retained);

        guard.client_mut().disconnect().await
    }

    pub(crate) fn begin_shutdown(&self) {
        self.initial_connect_cancellation.load().cancel();
        self.signal.store(true, Ordering::Release);
    }

    /// Receives the next message from the handler, or `None` if the receiver
    /// has been taken or the handler has shut down.
    pub async fn next_event(&mut self) -> Option<NautilusWsMessage> {
        if let Some(rx) = self.out_rx.as_mut() {
            rx.recv().await
        } else {
            None
        }
    }

    /// Takes the feed-handler task slot, leaving an empty slot behind.
    ///
    /// Used by callers that connect on a cloned client and want to await the
    /// inner handler task on a different instance during disconnect.
    #[must_use]
    pub(crate) fn take_task_slot(&mut self) -> TaskSlot<()> {
        std::mem::take(&mut self.task_handle)
    }

    /// Installs a feed-handler task slot previously obtained from [`Self::take_task_slot`].
    pub(crate) fn set_task_slot(&mut self, slot: TaskSlot<()>) {
        assert!(self.task_handle.is_none(), "task slot is already occupied");
        self.task_handle = slot;
    }

    /// Subscribe to L2 order-book updates for an instrument.
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered, the command
    /// cannot be queued, or the venue rejects the subscription.
    pub async fn subscribe_book(&self, instrument_id: InstrumentId) -> Result<(), LighterWsError> {
        let market_index = self.market_index_for(&instrument_id)?;

        let session = self
            .send_cmd_session(HandlerCommand::SetBookDeltasSub {
                market_index,
                subscribed: true,
            })
            .await?;

        if let Err(e) = self
            .subscribe_order_book_stream(market_index, &session)
            .await
        {
            let _ = self
                .send_cmd_on_session(
                    HandlerCommand::SetBookDeltasSub {
                        market_index,
                        subscribed: false,
                    },
                    &session,
                )
                .await;
            return Err(e);
        }

        Ok(())
    }

    /// Unsubscribe from L2 order-book updates.
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered or the command
    /// cannot be queued.
    pub async fn unsubscribe_book(
        &self,
        instrument_id: InstrumentId,
    ) -> Result<(), LighterWsError> {
        let market_index = self.market_index_for(&instrument_id)?;
        self.send_cmd(HandlerCommand::SetBookDeltasSub {
            market_index,
            subscribed: false,
        })
        .await?;
        self.unsubscribe_order_book_stream(market_index).await
    }

    /// Subscribe to depth-10 snapshots derived from the same `order_book`
    /// stream as [`Self::subscribe_book`].
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered, the command
    /// cannot be queued, or the venue rejects the subscription.
    pub async fn subscribe_book_depth(
        &self,
        instrument_id: InstrumentId,
    ) -> Result<(), LighterWsError> {
        let market_index = self.market_index_for(&instrument_id)?;

        let session = self
            .send_cmd_session(HandlerCommand::SetDepthSub {
                market_index,
                subscribed: true,
            })
            .await?;

        if let Err(e) = self
            .subscribe_order_book_stream(market_index, &session)
            .await
        {
            let _ = self
                .send_cmd_on_session(
                    HandlerCommand::SetDepthSub {
                        market_index,
                        subscribed: false,
                    },
                    &session,
                )
                .await;
            return Err(e);
        }

        Ok(())
    }

    /// Unsubscribe from depth-10 snapshots.
    ///
    /// Clears the depth-10 emission flag without tearing down the underlying
    /// `order_book` stream so any active deltas subscriber keeps receiving
    /// updates.
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered or the command
    /// cannot be queued.
    pub async fn unsubscribe_book_depth(
        &self,
        instrument_id: InstrumentId,
    ) -> Result<(), LighterWsError> {
        let market_index = self.market_index_for(&instrument_id)?;
        self.send_cmd(HandlerCommand::SetDepthSub {
            market_index,
            subscribed: false,
        })
        .await?;
        self.unsubscribe_order_book_stream(market_index).await
    }

    /// Subscribe to ticker (best bid/offer) updates.
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered, the command
    /// cannot be queued, or the venue rejects the subscription.
    pub async fn subscribe_quotes(
        &self,
        instrument_id: InstrumentId,
    ) -> Result<(), LighterWsError> {
        let market_index = self.market_index_for(&instrument_id)?;
        self.send_subscribe(LighterWsChannel::Ticker(market_index), None)
            .await
    }

    /// Unsubscribe from ticker updates.
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered or the command
    /// cannot be queued.
    pub async fn unsubscribe_quotes(
        &self,
        instrument_id: InstrumentId,
    ) -> Result<(), LighterWsError> {
        let market_index = self.market_index_for(&instrument_id)?;
        self.send_unsubscribe(LighterWsChannel::Ticker(market_index))
            .await
    }

    /// Subscribe to trade updates.
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered, the command
    /// cannot be queued, or the venue rejects the subscription.
    pub async fn subscribe_trades(
        &self,
        instrument_id: InstrumentId,
    ) -> Result<(), LighterWsError> {
        let market_index = self.market_index_for(&instrument_id)?;
        self.send_subscribe(LighterWsChannel::Trade(market_index), None)
            .await
    }

    /// Unsubscribe from trade updates.
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered or the command
    /// cannot be queued.
    pub async fn unsubscribe_trades(
        &self,
        instrument_id: InstrumentId,
    ) -> Result<(), LighterWsError> {
        let market_index = self.market_index_for(&instrument_id)?;
        self.send_unsubscribe(LighterWsChannel::Trade(market_index))
            .await
    }

    /// Subscribe to the `candle/{market_id}/{resolution}` stream for an
    /// instrument and resolution.
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered, the resolution
    /// is not offered on the WebSocket stream, the command cannot be queued,
    /// or the venue rejects the subscription.
    pub async fn subscribe_candles(
        &self,
        instrument_id: InstrumentId,
        resolution: LighterCandleResolution,
    ) -> Result<(), LighterWsError> {
        if !resolution.is_ws_streamable() {
            return Err(LighterWsError::Client(format!(
                "resolution {resolution:?} is not offered on the Lighter candle WebSocket stream",
            )));
        }
        let market_index = self.market_index_for(&instrument_id)?;
        self.send_subscribe(
            LighterWsChannel::Candle {
                market_index,
                resolution,
            },
            None,
        )
        .await
    }

    /// Unsubscribe from a candle stream.
    ///
    /// # Errors
    ///
    /// Returns an error if the instrument is not registered or the command
    /// cannot be queued.
    pub async fn unsubscribe_candles(
        &self,
        instrument_id: InstrumentId,
        resolution: LighterCandleResolution,
    ) -> Result<(), LighterWsError> {
        let market_index = self.market_index_for(&instrument_id)?;
        self.send_unsubscribe(LighterWsChannel::Candle {
            market_index,
            resolution,
        })
        .await
    }

    /// Subscribe to a market-stats stream covering all markets or a single
    /// market index.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued or the venue rejects
    /// the subscription.
    pub async fn subscribe_market_stats(
        &self,
        selection: LighterMarketSelection,
    ) -> Result<(), LighterWsError> {
        self.send_subscribe(LighterWsChannel::MarketStats(selection), None)
            .await
    }

    /// Unsubscribe from a market-stats stream.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued.
    pub async fn unsubscribe_market_stats(
        &self,
        selection: LighterMarketSelection,
    ) -> Result<(), LighterWsError> {
        self.send_unsubscribe(LighterWsChannel::MarketStats(selection))
            .await
    }

    /// Subscribe to a spot market-stats stream covering all spot markets or a
    /// single spot market index.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued or the venue rejects
    /// the subscription.
    pub async fn subscribe_spot_market_stats(
        &self,
        selection: LighterMarketSelection,
    ) -> Result<(), LighterWsError> {
        self.send_subscribe(LighterWsChannel::SpotMarketStats(selection), None)
            .await
    }

    /// Unsubscribe from a spot market-stats stream.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued.
    pub async fn unsubscribe_spot_market_stats(
        &self,
        selection: LighterMarketSelection,
    ) -> Result<(), LighterWsError> {
        self.send_unsubscribe(LighterWsChannel::SpotMarketStats(selection))
            .await
    }

    /// Subscribe to the chain-height stream.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued or the venue rejects
    /// the subscription.
    pub async fn subscribe_height(&self) -> Result<(), LighterWsError> {
        self.send_subscribe(LighterWsChannel::Height, None).await
    }

    /// Unsubscribe from the chain-height stream.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued.
    pub async fn unsubscribe_height(&self) -> Result<(), LighterWsError> {
        self.send_unsubscribe(LighterWsChannel::Height).await
    }

    /// Provides the execution context the feed handler stamps onto reports
    /// parsed from `account_*` frames.
    ///
    /// Without this context account frames fall back to
    /// [`NautilusWsMessage::Raw`]; once it is set the handler emits typed
    /// [`crate::websocket::messages::ExecutionReport`] and
    /// [`crate::websocket::messages::NautilusWsMessage::AccountState`]
    /// messages stamped with `account_id`. The `account_index` is used by
    /// the fill parser to determine which side of each account-trade frame
    /// the configured account took.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued.
    pub async fn set_execution_context(
        &self,
        account_id: AccountId,
        account_index: i64,
    ) -> Result<(), LighterWsError> {
        self.send_cmd(HandlerCommand::SetExecutionContext {
            account_id,
            account_index,
        })
        .await
    }

    /// Subscribe to a private account channel using a venue auth token.
    ///
    /// The auth token must be a valid Lighter L2 auth signature; see the
    /// `signing` module for token construction. Re-issuing this with a fresh
    /// token (as the execution client's auth-token rotation does) overwrites
    /// the stored reconnect-replay token, keeping it within the venue TTL.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued or the venue rejects
    /// the subscription.
    pub async fn subscribe_account(
        &self,
        channel: LighterWsChannel,
        auth_token: SecretString,
    ) -> Result<(), LighterWsError> {
        self.send_subscribe(channel, Some(auth_token)).await
    }

    /// Unsubscribe from a private account channel.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued.
    pub async fn unsubscribe_account(
        &self,
        channel: LighterWsChannel,
    ) -> Result<(), LighterWsError> {
        self.send_unsubscribe(channel).await
    }

    /// Dispatch a signed L2 transaction over the WebSocket.
    ///
    /// `tx_type` is the venue's [`crate::common::enums::LighterTxType`]
    /// discriminant; `tx_info` is the JSON body produced by the matching
    /// [`crate::signing::tx::TxInfoJson`] renderer. The venue confirms
    /// acceptance via the `account_*` streams.
    ///
    /// # Errors
    ///
    /// Returns an error if the command cannot be queued or the handler cannot
    /// report whether it handed the frame to the network writer.
    pub async fn send_tx(
        &self,
        tx_type: u8,
        tx_info: Box<serde_json::value::RawValue>,
    ) -> Result<(), LighterWsError> {
        self.send_tx_on_connection(tx_type, tx_info, self.connection_epoch())
            .await
    }

    pub(crate) async fn send_tx_on_connection(
        &self,
        tx_type: u8,
        tx_info: Box<serde_json::value::RawValue>,
        connection_epoch: u64,
    ) -> Result<(), LighterWsError> {
        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
        self.send_cmd(HandlerCommand::SendTx {
            tx_type,
            tx_info,
            connection_epoch,
            response_tx,
        })
        .await?;

        response_rx.await.map_err(|e| {
            LighterWsError::SendTxOutcomeUnknown(format!(
                "handler dropped sendTx result after accepting the command: {e}",
            ))
        })?
    }

    pub(crate) async fn send_tx_batch_on_connection(
        &self,
        data: super::messages::LighterWsSendTxBatch,
        connection_epoch: u64,
    ) -> Result<(), LighterWsError> {
        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
        self.send_cmd(HandlerCommand::SendTxBatch {
            data,
            connection_epoch,
            response_tx,
        })
        .await?;

        response_rx.await.map_err(|e| {
            LighterWsError::SendTxOutcomeUnknown(format!(
                "handler dropped sendTxBatch result after accepting the command: {e}",
            ))
        })?
    }

    async fn send_subscribe(
        &self,
        channel: LighterWsChannel,
        auth: Option<SecretString>,
    ) -> Result<(), LighterWsError> {
        self.send_subscribe_on_session(channel, auth, None).await
    }

    async fn send_subscribe_on_session(
        &self,
        channel: LighterWsChannel,
        auth: Option<SecretString>,
        expected_session: Option<&tokio::sync::mpsc::UnboundedSender<HandlerCommand>>,
    ) -> Result<(), LighterWsError> {
        let topic = channel.topic_key();
        let sender = self.cmd_tx.read().await;

        if expected_session.is_some_and(|session| !sender.same_channel(session)) {
            return Err(LighterWsError::Client(format!(
                "WebSocket session changed before subscription for {topic}",
            )));
        }

        if matches!(channel, LighterWsChannel::OrderBook(_))
            && !self.subscriptions.add_reference(&topic)
        {
            return Ok(());
        }

        let generation = self
            .next_subscription_generation
            .fetch_add(1, Ordering::Relaxed);
        let previous = self.subscription_args.insert(
            topic.clone(),
            SubscriptionArgs {
                channel: channel.clone(),
                auth: auth.clone(),
                generation,
            },
        );
        let (response_tx, response_rx) = tokio::sync::oneshot::channel();

        if let Err(e) = send_command(
            &sender,
            HandlerCommand::Subscribe {
                channel: channel.clone(),
                auth,
                response_tx: Some(response_tx),
            },
        ) {
            if matches!(channel, LighterWsChannel::OrderBook(_)) {
                self.remove_subscription_args(&channel, generation);
            } else {
                self.restore_subscription_args(&topic, generation, previous);
            }

            return Err(e);
        }

        let session = sender.clone();
        drop(sender);

        let result = match response_rx.await {
            Ok(result) => result.map_err(LighterWsError::Client),
            Err(e) => Err(LighterWsError::Client(format!(
                "handler dropped subscription result for {topic}: {e}",
            ))),
        };

        let sender = self.cmd_tx.read().await;

        if result.is_err() && sender.same_channel(&session) {
            self.remove_subscription_args(&channel, generation);
        }

        result
    }

    fn restore_subscription_args(
        &self,
        topic: &str,
        generation: u64,
        previous: Option<SubscriptionArgs>,
    ) {
        let Entry::Occupied(mut entry) = self.subscription_args.entry(topic.to_string()) else {
            return;
        };

        if entry.get().generation != generation {
            return;
        }

        if let Some(previous) = previous {
            entry.insert(previous);
        } else {
            entry.remove();
        }
    }

    fn remove_subscription_args(&self, channel: &LighterWsChannel, generation: u64) {
        let topic = channel.topic_key();
        // Remove the failed caller before deciding whether another book consumer needs replay.
        if matches!(channel, LighterWsChannel::OrderBook(_))
            && !self.subscriptions.remove_reference(&topic)
        {
            return;
        }

        let Entry::Occupied(entry) = self.subscription_args.entry(topic) else {
            return;
        };

        if entry.get().generation == generation {
            entry.remove();
        }
    }

    async fn send_unsubscribe(&self, channel: LighterWsChannel) -> Result<(), LighterWsError> {
        let topic = channel.topic_key();
        self.send_cmd(HandlerCommand::Unsubscribe { channel })
            .await?;
        self.subscription_args.remove(&topic);
        Ok(())
    }

    async fn subscribe_order_book_stream(
        &self,
        market_index: i64,
        session: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
    ) -> Result<(), LighterWsError> {
        self.send_subscribe_on_session(
            LighterWsChannel::OrderBook(market_index),
            None,
            Some(session),
        )
        .await
    }

    async fn unsubscribe_order_book_stream(&self, market_index: i64) -> Result<(), LighterWsError> {
        let channel = LighterWsChannel::OrderBook(market_index);
        let topic = channel.topic_key();

        if !self.subscriptions.remove_reference(topic.as_str()) {
            return Ok(());
        }

        if let Err(e) = self.send_unsubscribe(channel).await {
            self.subscriptions.add_reference(topic.as_str());
            return Err(e);
        }

        Ok(())
    }

    async fn send_cmd(&self, cmd: HandlerCommand) -> Result<(), LighterWsError> {
        let sender = self.cmd_tx.read().await;
        send_command(&sender, cmd)
    }

    async fn send_cmd_session(
        &self,
        cmd: HandlerCommand,
    ) -> Result<tokio::sync::mpsc::UnboundedSender<HandlerCommand>, LighterWsError> {
        let sender = self.cmd_tx.read().await;
        send_command(&sender, cmd)?;
        Ok(sender.clone())
    }

    async fn send_cmd_on_session(
        &self,
        cmd: HandlerCommand,
        session: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
    ) -> Result<(), LighterWsError> {
        let sender = self.cmd_tx.read().await;
        if sender.same_channel(session) {
            send_command(&sender, cmd)?;
        }

        Ok(())
    }

    fn market_index_for(&self, instrument_id: &InstrumentId) -> Result<i64, LighterWsError> {
        self.registry.market_index(instrument_id).ok_or_else(|| {
            LighterWsError::Client(format!(
                "no Lighter market_index registered for instrument: {instrument_id}"
            ))
        })
    }
}

impl Drop for LighterWebSocketClient {
    fn drop(&mut self) {
        if self.task_handle.is_none() {
            return;
        }

        self.connection_generation.fetch_add(1, Ordering::AcqRel);
        self.initial_connect_cancellation.load().cancel();
        self.signal.store(true, Ordering::Release);

        if let Some(handle) = self.task_handle.as_ref() {
            handle.abort();
        }

        if let Some(control) = &self.socket_control {
            control.deregister();
        }
    }
}

fn send_command(
    sender: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
    cmd: HandlerCommand,
) -> Result<(), LighterWsError> {
    sender
        .send(cmd)
        .map_err(|e| LighterWsError::Client(format!("handler unavailable: {e}")))
}

#[cfg(test)]
mod tests {
    use futures_util::{SinkExt, StreamExt};
    use nautilus_common::testing::wait_until_async;
    use nautilus_core::{UnixNanos, string::secret::REDACTED};
    use nautilus_model::{
        data::TradeTick,
        enums::AggressorSide,
        identifiers::{Symbol, TradeId},
        instruments::CryptoPerpetual,
        types::{Currency, Price, Quantity},
    };
    use nautilus_testkit::connection::race_command_with_publication;
    use rstest::rstest;

    use super::*;
    use crate::common::{
        consts::LIGHTER_VENUE,
        enums::{LighterProductType, LighterTxType},
    };

    impl LighterWebSocketClient {
        pub(crate) async fn drop_next_send_tx_result_for_test(&self) {
            let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
            *self.cmd_tx.write().await = cmd_tx;

            get_runtime().spawn(async move {
                if let Some(
                    HandlerCommand::SendTx { response_tx, .. }
                    | HandlerCommand::SendTxBatch { response_tx, .. },
                ) = cmd_rx.recv().await
                {
                    cmd_rx.close();
                    drop(response_tx);
                }
            });
        }
    }

    fn registry_with(
        market_index: i64,
        symbol: &str,
        product: LighterProductType,
    ) -> Arc<MarketRegistry> {
        let registry = Arc::new(MarketRegistry::new());
        registry.insert(market_index, symbol, product);
        registry
    }

    #[tokio::test]
    async fn command_during_sender_publication_reaches_initialized_handler() {
        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
        let addr = listener.local_addr().unwrap();
        let (received_tx, mut received_rx) = tokio::sync::mpsc::unbounded_channel();

        let server = tokio::spawn(async move {
            let (socket, _) = listener.accept().await.unwrap();
            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
            while let Some(Ok(message)) = socket.next().await {
                if let tokio_tungstenite::tungstenite::Message::Text(text) = message {
                    let payload: serde_json::Value = serde_json::from_str(text.as_str()).unwrap();
                    let _ = received_tx.send(payload);
                    socket
                        .send(tokio_tungstenite::tungstenite::Message::Text(
                            include_str!("../../test_data/ws_trade_update.json").into(),
                        ))
                        .await
                        .unwrap();
                }
            }
        });

        let url = format!("ws://{addr}");

        let mut ws_client = LighterWebSocketClient::new(
            Some(url),
            LighterEnvironment::Testnet,
            registry_with(0, "ETH", LighterProductType::Perp),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        ws_client
            .connection_epoch
            .store(Arc::new(AtomicU64::new(99)));
        let observer = ws_client.clone();
        let published = ws_client.clone();
        let instrument_id = InstrumentId::from("ETH-USD-PERP.LIGHTER");

        let command = HandlerCommand::Subscribe {
            channel: LighterWsChannel::Trade(0),
            auth: None,
            response_tx: None,
        };

        race_command_with_publication(
            Arc::clone(&ws_client.cmd_tx),
            ws_client.connect(),
            move || {
                assert!(published.is_active());
                assert_eq!(published.connection_epoch(), 0);
                command
            },
            || {
                observer.cache_instrument(0, stub_instrument(instrument_id));
                observer.is_active()
            },
        )
        .await
        .expect("connect failed");

        let received = tokio::time::timeout(Duration::from_secs(2), received_rx.recv()).await;
        let event = tokio::time::timeout(Duration::from_secs(2), ws_client.next_event()).await;
        ws_client.disconnect().await.expect("disconnect failed");
        server.abort();

        let payload = received
            .expect("startup command must reach server")
            .unwrap();
        assert_eq!(
            payload,
            serde_json::json!({
                "type": "subscribe",
                "channel": "trade/0"
            })
        );

        let NautilusWsMessage::Trades(ticks) = event
            .expect("startup cache must reach handler")
            .expect("trade event")
        else {
            panic!("expected parsed trades");
        };

        assert_eq!(ticks.len(), 1);
        assert_eq!(
            ticks[0],
            TradeTick::new(
                instrument_id,
                Price::from("2181.83"),
                Quantity::from("0.1336"),
                AggressorSide::Sell,
                TradeId::from("16164557907"),
                UnixNanos::from(1_773_854_156_654_000_000),
                ticks[0].ts_init,
            )
        );
    }

    #[rstest]
    #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
    #[case(
        "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
        "wss://stream.example.com:9443"
    )]
    fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
        let client = LighterWebSocketClient::new(
            Some(url.to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let debug = format!("{client:?}");
        assert!(debug.contains(&format!("url: {origin:?}")));
        assert!(!debug.contains("SYN-"));
        assert_eq!(client.url.expose(), url);
    }

    #[rstest]
    #[case::only_caller(1, false, false)]
    #[case::other_consumer(2, false, true)]
    #[case::other_consumer_leaves(2, true, false)]
    fn book_failed_initial_waiter_preserves_other_consumer_replay(
        #[case] references: usize,
        #[case] other_leaves: bool,
        #[case] retained: bool,
    ) {
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".into()),
            LighterEnvironment::Testnet,
            registry_with(0, "ETH", LighterProductType::Perp),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let topic = "order_book:0";
        for _ in 0..references {
            client.subscriptions.add_reference(topic);
        }

        client.subscription_args.insert(
            topic.into(),
            SubscriptionArgs {
                channel: LighterWsChannel::OrderBook(0),
                auth: None,
                generation: 7,
            },
        );

        if other_leaves {
            assert!(!client.subscriptions.remove_reference(topic));
        }

        client.remove_subscription_args(&LighterWsChannel::OrderBook(0), 7);

        assert_eq!(client.subscription_args.contains_key(topic), retained);
        assert_eq!(
            client.subscriptions.get_reference_count(topic),
            references - 1 - usize::from(other_leaves)
        );

        if retained {
            assert_eq!(client.subscription_args.get(topic).unwrap().generation, 7);
        }
    }

    #[rstest]
    fn debug_redacts_retained_auth_and_proxy_url() {
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            Some("http://user:proxy-secret@localhost".to_string()),
        );
        client.subscription_args.insert(
            "account_all:7".to_string(),
            SubscriptionArgs {
                channel: LighterWsChannel::AccountAll(7),
                auth: Some(SecretString::from("auth-token")),
                generation: 1,
            },
        );

        let debug = format!("{client:?}");

        assert!(debug.contains(REDACTED));
        assert!(!debug.contains("proxy-secret"));
        assert!(!debug.contains("auth-token"));
    }

    #[rstest]
    fn market_index_for_returns_registered_index() {
        let registry = registry_with(7, "ETH", LighterProductType::Perp);
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::clone(&registry),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let id = registry.instrument_id(7).expect("registered");
        assert_eq!(client.market_index_for(&id).unwrap(), 7);
    }

    #[rstest]
    fn market_index_for_unregistered_returns_error() {
        let registry = Arc::new(MarketRegistry::new());
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            registry,
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let id = InstrumentId::new(Symbol::from_str_unchecked("UNKNOWN-PERP"), *LIGHTER_VENUE);
        assert!(client.market_index_for(&id).is_err());
    }

    #[rstest]
    fn cache_instrument_populates_lookup() {
        let registry = registry_with(0, "ETH", LighterProductType::Perp);
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::clone(&registry),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let id = registry.instrument_id(0).expect("registered");
        let instrument = stub_instrument(id);
        client.cache_instrument(0, instrument);
        assert!(client.instruments_cache().contains_key(&0));
    }

    #[tokio::test]
    async fn wait_until_active_uses_configured_timeout() {
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            0,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );

        let error = client
            .wait_until_active()
            .await
            .expect_err("inactive client should time out");

        assert!(error.to_string().contains("timeout after 0 seconds"));
    }

    #[rstest]
    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn disconnect_awaits_handler_after_timeout_abort() {
        struct NotifyOnDrop {
            tx: Option<tokio::sync::oneshot::Sender<()>>,
        }

        impl Drop for NotifyOnDrop {
            fn drop(&mut self) {
                if let Some(tx) = self.tx.take() {
                    let _ = tx.send(());
                }
            }
        }

        let mut client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
        let (drop_tx, mut drop_rx) = tokio::sync::oneshot::channel();
        let (release_tx, release_rx) = std::sync::mpsc::channel();

        client.task_handle.insert(get_runtime().spawn(async move {
            let _notify = NotifyOnDrop { tx: Some(drop_tx) };
            let _ = started_tx.send(());
            let _ = release_rx.recv();
            std::future::pending::<()>().await;
        }));

        tokio::time::timeout(Duration::from_secs(2), started_rx)
            .await
            .expect("handler startup deadline")
            .expect("handler task started");

        let disconnect = client.disconnect();
        tokio::pin!(disconnect);
        assert!(
            tokio::time::timeout(
                DISCONNECT_TIMEOUT + Duration::from_millis(250),
                &mut disconnect,
            )
            .await
            .is_err(),
            "disconnect must await the blocked handler after aborting it",
        );
        assert_eq!(
            drop_rx.try_recv(),
            Err(tokio::sync::oneshot::error::TryRecvError::Empty)
        );
        release_tx.send(()).expect("release blocked handler");
        tokio::time::timeout(Duration::from_secs(2), disconnect)
            .await
            .expect("handler must finish after release")
            .expect("disconnect");

        assert_eq!(drop_rx.try_recv(), Ok(()));
    }

    #[tokio::test]
    async fn drop_clone_does_not_cancel_handler() {
        let mut client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        client
            .task_handle
            .insert(get_runtime().spawn(std::future::pending()));
        let clone = client.clone();

        drop(clone);

        assert!(!client.signal.load(Ordering::Acquire));
        assert!(
            !client
                .task_handle
                .as_ref()
                .expect("handler task")
                .is_finished()
        );
    }

    #[tokio::test]
    async fn cancelled_moved_disconnect_retains_handler_task() {
        let mut client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        client
            .task_handle
            .insert(get_runtime().spawn(std::future::pending()));
        let signal = Arc::clone(&client.signal);
        let retained = Arc::new(RetainedTaskSlot::new());
        let disconnect_future = client.disconnect_with_task_retention(Arc::clone(&retained));

        let disconnect = get_runtime().spawn(disconnect_future);

        wait_until_async(
            || async { signal.load(Ordering::Acquire) },
            Duration::from_secs(2),
        )
        .await;
        disconnect.abort();
        let _ = disconnect.await;

        assert!(!retained.is_empty());
        retained.0.abort();
        retained.finish().await.expect("retained handler shutdown");
        assert!(retained.is_empty());
    }

    #[tokio::test]
    async fn send_tx_reports_unknown_outcome_when_handler_result_is_dropped() {
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        client.drop_next_send_tx_result_for_test().await;
        let tx_info = serde_json::value::RawValue::from_string("{}".to_string()).unwrap();

        let error = client
            .send_tx(LighterTxType::CreateOrder as u8, tx_info)
            .await
            .expect_err("dropped handler result must be ambiguous");

        assert!(matches!(error, LighterWsError::SendTxOutcomeUnknown(_)));
    }

    #[tokio::test]
    async fn subscription_registration_waits_for_session_publication() {
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let mut publication = client.cmd_tx.write().await;
        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
        let mut subscribe =
            Box::pin(client.subscribe_market_stats(LighterMarketSelection::Market(0)));

        assert!(futures_util::poll!(subscribe.as_mut()).is_pending());
        assert!(client.subscription_args.is_empty());

        client.subscriptions.clear();
        client.subscription_args.clear();
        *publication = cmd_tx;
        drop(publication);

        assert!(futures_util::poll!(subscribe.as_mut()).is_pending());
        let command = cmd_rx.try_recv().expect("new-session subscribe command");

        let HandlerCommand::Subscribe {
            channel,
            auth,
            response_tx: Some(response_tx),
        } = command
        else {
            panic!("expected subscribe command with venue result sender");
        };

        response_tx
            .send(Ok(()))
            .expect("subscription result receiver");
        subscribe.await.expect("new-session subscription");

        assert_eq!(
            channel,
            LighterWsChannel::MarketStats(LighterMarketSelection::Market(0))
        );
        assert!(auth.is_none());
        assert_eq!(client.subscription_args.len(), 1);
        let args = client
            .subscription_args
            .get("market_stats:0")
            .expect("replay arguments");
        assert_eq!(args.channel, channel);
        assert!(args.auth.is_none());
        assert_eq!(args.generation, 1);
    }

    #[rstest]
    #[case::deltas(false)]
    #[case::depth(true)]
    #[tokio::test]
    async fn book_registration_rejects_previous_session(#[case] depth: bool) {
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let market_index = 7;
        let generation = 42;
        let channel = LighterWsChannel::OrderBook(market_index);
        let topic = channel.topic_key();

        let consumer = |subscribed| {
            if depth {
                HandlerCommand::SetDepthSub {
                    market_index,
                    subscribed,
                }
            } else {
                HandlerCommand::SetBookDeltasSub {
                    market_index,
                    subscribed,
                }
            }
        };

        let (previous_tx, mut previous_rx) = tokio::sync::mpsc::unbounded_channel();
        *client.cmd_tx.write().await = previous_tx;
        let session = client
            .send_cmd_session(consumer(true))
            .await
            .expect("previous-session consumer registration");
        let (current_tx, mut current_rx) = tokio::sync::mpsc::unbounded_channel();
        *client.cmd_tx.write().await = current_tx;
        client.subscriptions.add_reference(&topic);
        client.subscriptions.mark_subscribe(&topic);
        client.subscriptions.confirm_subscribe(&topic);
        client
            .next_subscription_generation
            .store(generation + 1, Ordering::Relaxed);
        client.subscription_args.insert(
            topic.clone(),
            SubscriptionArgs {
                channel: channel.clone(),
                auth: None,
                generation,
            },
        );

        let error = client
            .subscribe_order_book_stream(market_index, &session)
            .await
            .expect_err("previous-session book registration must fail");
        client
            .send_cmd_on_session(consumer(false), &session)
            .await
            .expect("skip previous-session rollback");

        let LighterWsError::Client(message) = error else {
            panic!("expected session mismatch client error");
        };

        assert_eq!(
            message,
            format!("WebSocket session changed before subscription for {topic}")
        );
        assert_eq!(client.subscriptions.get_reference_count(&topic), 1);
        assert_eq!(client.subscription_count(), 1);
        assert_eq!(client.subscription_args.len(), 1);
        assert_eq!(
            client.next_subscription_generation.load(Ordering::Relaxed),
            generation + 1
        );
        let args = client
            .subscription_args
            .get(&topic)
            .expect("current-session replay args");
        assert_eq!(args.channel, channel);
        assert!(args.auth.is_none());
        assert_eq!(args.generation, generation);
        let command = previous_rx
            .try_recv()
            .expect("previous-session consumer flag");

        match (depth, command) {
            (
                true,
                HandlerCommand::SetDepthSub {
                    market_index: actual,
                    subscribed,
                },
            )
            | (
                false,
                HandlerCommand::SetBookDeltasSub {
                    market_index: actual,
                    subscribed,
                },
            ) => {
                assert_eq!(actual, market_index);
                assert!(subscribed);
            }
            (_, command) => panic!("unexpected previous-session command: {command:?}"),
        }

        assert!(matches!(
            previous_rx.try_recv(),
            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
        ));
        assert!(matches!(
            current_rx.try_recv(),
            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
        ));
    }

    #[tokio::test]
    async fn subscribe_waits_for_venue_result_and_removes_failed_generation() {
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
        *client.cmd_tx.write().await = cmd_tx;

        let subscribe_client = client.clone();
        let subscribe = get_runtime().spawn(async move {
            subscribe_client
                .subscribe_market_stats(LighterMarketSelection::Market(0))
                .await
        });

        let command = cmd_rx.recv().await.expect("subscribe command");
        let HandlerCommand::Subscribe {
            response_tx: Some(response_tx),
            ..
        } = command
        else {
            panic!("expected subscribe command with venue result sender");
        };

        assert!(!subscribe.is_finished());
        assert!(client.subscription_args.contains_key("market_stats:0"));

        response_tx
            .send(Err("venue rejected subscription".to_string()))
            .expect("subscription result receiver");
        let error = subscribe
            .await
            .expect("subscribe task")
            .expect_err("failed venue open must fail the caller");

        assert!(error.to_string().contains("venue rejected subscription"));
        assert!(!client.subscription_args.contains_key("market_stats:0"));
    }

    #[tokio::test]
    async fn failed_older_subscribe_does_not_remove_newer_generation() {
        let client = LighterWebSocketClient::new(
            Some("wss://example/test".to_string()),
            LighterEnvironment::Testnet,
            Arc::new(MarketRegistry::new()),
            TransportBackend::default(),
            30,
            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
            None,
        );
        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
        *client.cmd_tx.write().await = cmd_tx;

        let older_client = client.clone();
        let older = get_runtime().spawn(async move {
            older_client
                .subscribe_market_stats(LighterMarketSelection::Market(0))
                .await
        });
        let HandlerCommand::Subscribe {
            response_tx: Some(older_response),
            ..
        } = cmd_rx.recv().await.expect("older subscribe command")
        else {
            panic!("expected older subscribe command with venue result sender");
        };

        let newer_client = client.clone();
        let newer = get_runtime().spawn(async move {
            newer_client
                .subscribe_market_stats(LighterMarketSelection::Market(0))
                .await
        });
        let HandlerCommand::Subscribe {
            response_tx: Some(newer_response),
            ..
        } = cmd_rx.recv().await.expect("newer subscribe command")
        else {
            panic!("expected newer subscribe command with venue result sender");
        };

        older_response
            .send(Err("older generation failed".to_string()))
            .expect("older result receiver");
        older
            .await
            .expect("older subscribe task")
            .expect_err("older generation must fail");
        assert!(client.subscription_args.contains_key("market_stats:0"));

        newer_response
            .send(Err("newer generation failed".to_string()))
            .expect("newer result receiver");
        newer
            .await
            .expect("newer subscribe task")
            .expect_err("newer generation must fail");
        assert!(!client.subscription_args.contains_key("market_stats:0"));
    }

    fn stub_instrument(id: InstrumentId) -> InstrumentAny {
        InstrumentAny::CryptoPerpetual(
            CryptoPerpetual::builder()
                .instrument_id(id)
                .raw_symbol(id.symbol)
                .base_currency(Currency::from("ETH"))
                .quote_currency(Currency::from("USDC"))
                .settlement_currency(Currency::from("USDC"))
                .is_inverse(false)
                .price_precision(2)
                .size_precision(4)
                .price_increment(Price::from("0.01"))
                .size_increment(Quantity::from("0.0001"))
                .ts_event(UnixNanos::default())
                .ts_init(UnixNanos::default())
                .build()
                .unwrap(),
        )
    }
}