nautilus_lighter/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//! Outer WebSocket client orchestrating connection lifecycle and subscriptions.
17
18use std::{
19    fmt::Debug,
20    num::NonZeroU32,
21    sync::{
22        Arc,
23        atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
24    },
25    time::Duration,
26};
27
28use arc_swap::ArcSwap;
29use dashmap::{DashMap, mapref::entry::Entry};
30#[cfg(test)]
31use nautilus_common::live::get_runtime;
32use nautilus_core::string::secret::{SecretString, redact_option};
33#[cfg(test)]
34use nautilus_live::book::DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS;
35use nautilus_live::{
36    SocketControl,
37    task::{SharedTaskSlot, TaskJoinOutcome, TaskSlot, finish_task},
38};
39use nautilus_model::{
40    identifiers::{AccountId, InstrumentId},
41    instruments::InstrumentAny,
42};
43use nautilus_network::{
44    SocketStateSink,
45    http::create_standard_nautilus_headers,
46    mode::ConnectionMode,
47    websocket::{
48        InitialConnectRetryPolicy, SubscriptionState, TransportBackend, WebSocketClient,
49        WebSocketConfig, channel_epoch_message_handler,
50    },
51};
52use tokio_util::sync::CancellationToken;
53
54use crate::{
55    common::{
56        consts::{
57            DISCONNECT_TIMEOUT, HEARTBEAT_INTERVAL, HEARTBEAT_TIMEOUT, RECONNECT_BASE_BACKOFF,
58            RECONNECT_MAX_BACKOFF,
59        },
60        enums::{LighterCandleResolution, LighterEnvironment},
61        rate_limit::ws_message_rate_limiter,
62        symbol::MarketRegistry,
63        urls::lighter_ws_url,
64    },
65    websocket::{
66        error::LighterWsError,
67        handler::{FeedHandler, HandlerCommand},
68        messages::{LighterMarketSelection, LighterWsChannel, NautilusWsMessage},
69    },
70};
71
72const RECONNECT_JITTER_MS: u64 = 200;
73const RECONNECT_BACKOFF_FACTOR: f64 = 2.0;
74
75#[derive(Clone)]
76struct SubscriptionArgs {
77    channel: LighterWsChannel,
78    auth: Option<SecretString>,
79    generation: u64,
80}
81
82/// Outer Lighter WebSocket client.
83///
84/// Orchestrates the connection lifecycle and subscription bookkeeping for the
85/// Lighter streaming API. The inner feed handler runs on a dedicated tokio
86/// task and exclusively owns the underlying [`WebSocketClient`]; this outer
87/// type communicates with it through a command channel and consumes events
88/// over an unbounded mpsc.
89///
90/// Authenticated channels store their auth token in `subscription_args` and
91/// replay it verbatim on reconnect. That stored token stays valid because the
92/// execution client rotates it on a 6h cadence (inside the venue's 7h TTL) by
93/// re-issuing `subscribe_account`, so a reconnect never replays a connect-time,
94/// potentially-expired token.
95pub struct LighterWebSocketClient {
96    url: String,
97    connection_mode: Arc<ArcSwap<AtomicU8>>,
98    connection_epoch: Arc<ArcSwap<AtomicU64>>,
99    connection_lock: Arc<tokio::sync::Mutex<()>>,
100    connection_generation: Arc<AtomicU64>,
101    initial_connect_cancellation: Arc<ArcSwap<CancellationToken>>,
102    signal: Arc<AtomicBool>,
103    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
104    out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<NautilusWsMessage>>,
105    subscriptions: SubscriptionState,
106    subscription_args: Arc<DashMap<String, SubscriptionArgs>>,
107    next_subscription_generation: Arc<AtomicU64>,
108    instruments: Arc<DashMap<i64, InstrumentAny>>,
109    registry: Arc<MarketRegistry>,
110    task_handle: TaskSlot<()>,
111    transport_backend: TransportBackend,
112    ws_timeout_secs: u64,
113    book_snapshot_timeout: Duration,
114    proxy_url: Option<SecretString>,
115    socket_sink: Option<SocketStateSink>,
116    socket_control: Option<SocketControl>,
117}
118
119#[derive(Debug)]
120pub(crate) struct RetainedTaskSlot(SharedTaskSlot<()>);
121
122impl RetainedTaskSlot {
123    pub(crate) fn new() -> Self {
124        Self(SharedTaskSlot::new())
125    }
126
127    pub(crate) fn is_empty(&self) -> bool {
128        self.0.is_empty()
129    }
130
131    pub(crate) async fn finish(&self) -> Result<(), LighterWsError> {
132        let Some(outcome) = self.0.finish(DISCONNECT_TIMEOUT, DISCONNECT_TIMEOUT).await else {
133            return Ok(());
134        };
135
136        match outcome {
137            TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => Ok(()),
138            TaskJoinOutcome::Failed(error) => Err(LighterWsError::Client(format!(
139                "retained WebSocket handler task failed: {error}"
140            ))),
141            TaskJoinOutcome::Incomplete => Err(LighterWsError::Client(
142                "retained WebSocket handler task did not stop after abort".to_string(),
143            )),
144        }
145    }
146}
147
148pub(crate) struct TaskRetentionGuard {
149    client: Option<LighterWebSocketClient>,
150    retained: Arc<RetainedTaskSlot>,
151}
152
153impl TaskRetentionGuard {
154    pub(crate) fn new(client: LighterWebSocketClient, retained: Arc<RetainedTaskSlot>) -> Self {
155        Self {
156            client: Some(client),
157            retained,
158        }
159    }
160
161    pub(crate) fn client_mut(&mut self) -> &mut LighterWebSocketClient {
162        self.client.as_mut().expect("retention guard is armed")
163    }
164
165    pub(crate) fn disarm(mut self) -> LighterWebSocketClient {
166        self.client.take().expect("retention guard is armed")
167    }
168}
169
170impl Drop for TaskRetentionGuard {
171    fn drop(&mut self) {
172        let Some(client) = self.client.as_mut() else {
173            return;
174        };
175        client.begin_shutdown();
176        let slot = client.take_task_slot();
177        if slot.is_some() && self.retained.0.try_insert_slot(slot).is_err() {
178            log::error!("Lighter retained WebSocket task slot was already occupied");
179        }
180    }
181}
182
183impl Debug for LighterWebSocketClient {
184    /// Custom `Debug` that redacts the auth token in `subscription_args`.
185    ///
186    /// Authenticated channel subscriptions store the venue bearer token
187    /// alongside the channel for reconnect replay; deriving `Debug` would
188    /// otherwise print the token verbatim.
189    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
190        let subscription_topics: Vec<String> = self
191            .subscription_args
192            .iter()
193            .map(|entry| {
194                let args = entry.value();
195                format!(
196                    "topic={} channel={:?} authed={} generation={}",
197                    entry.key(),
198                    args.channel,
199                    args.auth.is_some(),
200                    args.generation,
201                )
202            })
203            .collect();
204
205        f.debug_struct(stringify!(LighterWebSocketClient))
206            .field("url", &self.url)
207            .field("is_active", &self.is_active())
208            .field("subscription_count", &self.subscriptions.len())
209            .field("subscription_args", &subscription_topics)
210            .field("instruments_len", &self.instruments.len())
211            .field("transport_backend", &self.transport_backend)
212            .field("ws_timeout_secs", &self.ws_timeout_secs)
213            .field("book_snapshot_timeout", &self.book_snapshot_timeout)
214            .field("proxy_url", &redact_option(self.proxy_url.as_ref()))
215            .finish_non_exhaustive()
216    }
217}
218
219impl Clone for LighterWebSocketClient {
220    fn clone(&self) -> Self {
221        Self {
222            url: self.url.clone(),
223            connection_mode: Arc::clone(&self.connection_mode),
224            connection_epoch: Arc::clone(&self.connection_epoch),
225            connection_lock: Arc::clone(&self.connection_lock),
226            connection_generation: Arc::clone(&self.connection_generation),
227            initial_connect_cancellation: Arc::clone(&self.initial_connect_cancellation),
228            signal: Arc::clone(&self.signal),
229            cmd_tx: Arc::clone(&self.cmd_tx),
230            out_rx: None,
231            subscriptions: self.subscriptions.clone(),
232            subscription_args: Arc::clone(&self.subscription_args),
233            next_subscription_generation: Arc::clone(&self.next_subscription_generation),
234            instruments: Arc::clone(&self.instruments),
235            registry: Arc::clone(&self.registry),
236            task_handle: TaskSlot::new(),
237            transport_backend: self.transport_backend,
238            ws_timeout_secs: self.ws_timeout_secs,
239            book_snapshot_timeout: self.book_snapshot_timeout,
240            proxy_url: self.proxy_url.clone(),
241            socket_sink: self.socket_sink.clone(),
242            socket_control: self.socket_control.clone(),
243        }
244    }
245}
246
247impl LighterWebSocketClient {
248    fn initial_connect_retry_policy() -> InitialConnectRetryPolicy {
249        InitialConnectRetryPolicy {
250            max_attempts: NonZeroU32::new(5).expect("initial connect attempts must be non-zero"),
251            delay_initial: Duration::from_millis(500),
252            delay_max: Duration::from_secs(5),
253            backoff_factor: 2.0,
254            jitter_ms: 250,
255        }
256    }
257
258    /// Creates a new client without connecting.
259    ///
260    /// `url` overrides the resolved environment URL when supplied.
261    #[must_use]
262    pub fn new(
263        url: Option<String>,
264        environment: LighterEnvironment,
265        registry: Arc<MarketRegistry>,
266        transport_backend: TransportBackend,
267        ws_timeout_secs: u64,
268        book_snapshot_timeout: Duration,
269        proxy_url: Option<String>,
270    ) -> Self {
271        let url = url.unwrap_or_else(|| lighter_ws_url(environment).to_string());
272        let connection_mode = Arc::new(ArcSwap::new(Arc::new(AtomicU8::new(
273            ConnectionMode::Closed as u8,
274        ))));
275        let connection_epoch = Arc::new(ArcSwap::new(Arc::new(AtomicU64::new(0))));
276
277        let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
278
279        Self {
280            url,
281            connection_mode,
282            connection_epoch,
283            connection_lock: Arc::new(tokio::sync::Mutex::new(())),
284            connection_generation: Arc::new(AtomicU64::new(0)),
285            initial_connect_cancellation: Arc::new(ArcSwap::from_pointee(CancellationToken::new())),
286            signal: Arc::new(AtomicBool::new(false)),
287            cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
288            out_rx: None,
289            subscriptions: SubscriptionState::new(':'),
290            subscription_args: Arc::new(DashMap::new()),
291            next_subscription_generation: Arc::new(AtomicU64::new(1)),
292            instruments: Arc::new(DashMap::new()),
293            registry,
294            task_handle: TaskSlot::new(),
295            transport_backend,
296            ws_timeout_secs,
297            book_snapshot_timeout,
298            proxy_url: proxy_url.map(SecretString::from),
299            socket_sink: None,
300            socket_control: None,
301        }
302    }
303
304    /// Configures socket state reporting for the underlying transport.
305    #[must_use]
306    pub fn with_state_sink(mut self, state_sink: SocketStateSink) -> Self {
307        self.socket_sink = Some(state_sink);
308        self
309    }
310
311    /// Configures socket state reporting and reconnect control.
312    #[must_use]
313    pub fn with_socket_control(mut self, control: SocketControl) -> Self {
314        self.socket_control = Some(control);
315        self
316    }
317
318    /// Returns the resolved WebSocket URL.
319    #[must_use]
320    pub fn url(&self) -> &str {
321        &self.url
322    }
323
324    /// Returns `true` when the underlying connection is active.
325    #[must_use]
326    pub fn is_active(&self) -> bool {
327        self.connection_mode.load().load(Ordering::Relaxed) == ConnectionMode::Active as u8
328    }
329
330    #[must_use]
331    pub(crate) fn connection_epoch(&self) -> u64 {
332        self.connection_epoch.load().load(Ordering::Acquire)
333    }
334
335    pub(crate) fn connection_epoch_atomic(&self) -> Arc<AtomicU64> {
336        self.connection_epoch.load_full()
337    }
338
339    /// Waits until the underlying connection reports active, or returns an
340    /// error after the configured WebSocket timeout.
341    ///
342    /// Polls [`Self::is_active`] every 10ms. Mirrors the documented
343    /// `wait_until_active` contract for adapter WebSocket clients in
344    /// `docs/developer_guide/adapters.md`.
345    ///
346    /// # Errors
347    ///
348    /// Returns [`LighterWsError::Client`] if the connection does not reach
349    /// the active state within the configured timeout.
350    pub async fn wait_until_active(&self) -> Result<(), LighterWsError> {
351        let timeout_secs = self.ws_timeout_secs;
352        let timeout = Duration::from_secs(timeout_secs);
353
354        tokio::time::timeout(timeout, async {
355            while !self.is_active() {
356                tokio::time::sleep(Duration::from_millis(10)).await;
357            }
358        })
359        .await
360        .map_err(|_| {
361            LighterWsError::Client(format!(
362                "WebSocket connection timeout after {timeout_secs} seconds"
363            ))
364        })
365    }
366
367    /// Returns the count of confirmed subscriptions.
368    #[must_use]
369    pub fn subscription_count(&self) -> usize {
370        self.subscriptions.len()
371    }
372
373    /// Returns a clone of the shared instrument cache.
374    #[must_use]
375    pub fn instruments_cache(&self) -> Arc<DashMap<i64, InstrumentAny>> {
376        Arc::clone(&self.instruments)
377    }
378
379    /// Caches a batch of instruments along with their venue `market_index`,
380    /// replaying them to the handler if a connection is already established.
381    pub fn cache_instruments(&self, instruments: Vec<(i64, InstrumentAny)>) {
382        self.instruments.clear();
383        for (market_index, instrument) in &instruments {
384            self.instruments.insert(*market_index, instrument.clone());
385        }
386        log::debug!(
387            "Lighter instrument cache initialized with {} instruments",
388            instruments.len()
389        );
390
391        if let Ok(cmd_tx) = self.cmd_tx.try_read() {
392            let _ = cmd_tx.send(HandlerCommand::InitializeInstruments(instruments));
393        }
394    }
395
396    /// Caches a single instrument and pushes it to the handler if connected.
397    pub fn cache_instrument(&self, market_index: i64, instrument: InstrumentAny) {
398        self.instruments.insert(market_index, instrument.clone());
399
400        if let Ok(cmd_tx) = self.cmd_tx.try_read() {
401            let _ = cmd_tx.send(HandlerCommand::UpdateInstrument {
402                market_index,
403                instrument,
404            });
405        }
406    }
407
408    /// Establishes the WebSocket connection and spawns the feed-handler task.
409    /// Classified transient failures retry within the configured WebSocket timeout.
410    ///
411    /// # Errors
412    ///
413    /// Returns an error if the connection fails permanently, exhausts its timeout, is cancelled,
414    /// or the handler cannot be initialized.
415    pub async fn connect(&mut self) -> anyhow::Result<()> {
416        self.connect_with_cancellation(CancellationToken::new())
417            .await
418    }
419
420    pub(crate) async fn connect_with_cancellation(
421        &mut self,
422        cancellation_token: CancellationToken,
423    ) -> anyhow::Result<()> {
424        let generation = self.connection_generation.load(Ordering::Acquire);
425        let _guard = self.connection_lock.lock().await;
426
427        anyhow::ensure!(
428            generation == self.connection_generation.load(Ordering::Acquire),
429            "Lighter WebSocket initial connection cancelled",
430        );
431
432        if self.is_active() {
433            log::warn!("Lighter WebSocket already connected");
434            return Ok(());
435        }
436
437        if let Some(outcome) = finish_task(
438            &mut self.task_handle,
439            DISCONNECT_TIMEOUT,
440            DISCONNECT_TIMEOUT,
441        )
442        .await
443        {
444            match outcome {
445                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
446                TaskJoinOutcome::Failed(error) => {
447                    anyhow::bail!("Lighter WebSocket handler failed: {error}");
448                }
449                TaskJoinOutcome::Incomplete => {
450                    anyhow::bail!("Lighter WebSocket handler did not stop after abort");
451                }
452            }
453        }
454
455        self.signal.store(false, Ordering::Release);
456        self.initial_connect_cancellation
457            .store(Arc::new(cancellation_token.clone()));
458
459        let (message_handler, raw_rx) = channel_epoch_message_handler();
460        let headers = create_standard_nautilus_headers();
461
462        let cfg = WebSocketConfig {
463            url: self.url.clone(),
464            headers,
465            heartbeat_interval_secs: Some(HEARTBEAT_INTERVAL.as_secs()),
466            heartbeat_payload: None,
467            connect_timeout_ms: Some(self.ws_timeout_secs.saturating_mul(1_000).max(1)),
468            reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
469            reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
470            reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
471            reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
472            reconnect_max_attempts: None,
473            heartbeat_timeout_secs: Some(HEARTBEAT_TIMEOUT.as_secs()),
474            idle_timeout_ms: None,
475            backend: self.transport_backend,
476            proxy_url: self
477                .proxy_url
478                .as_ref()
479                .map(|value| value.expose_secret().to_owned()),
480        };
481        let connect = WebSocketClient::epoch_builder()
482            .config(cfg)
483            .epoch_handler(message_handler)
484            .rate_limiter(ws_message_rate_limiter(&self.url))
485            .initial_connect_retry_policy(Self::initial_connect_retry_policy())
486            .cancellation_token(cancellation_token.clone())
487            .maybe_state_sink(
488                self.socket_control
489                    .as_ref()
490                    .map(SocketControl::sink)
491                    .or_else(|| self.socket_sink.clone()),
492            )
493            .connect();
494        let client =
495            match tokio::time::timeout(Duration::from_secs(self.ws_timeout_secs), connect).await {
496                Ok(result) => result?,
497                Err(_) => anyhow::bail!(
498                    "Lighter WebSocket initial connection timeout after {} seconds",
499                    self.ws_timeout_secs,
500                ),
501            };
502
503        if cancellation_token.is_cancelled()
504            || generation != self.connection_generation.load(Ordering::Acquire)
505        {
506            client.disconnect().await;
507
508            anyhow::bail!("Lighter WebSocket initial connection cancelled");
509        }
510
511        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
512        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<NautilusWsMessage>();
513
514        // Capture the connection-mode atomic before moving `client` into the
515        // SetClient command below.
516        let connection_mode_atomic = client.connection_mode_atomic();
517        let connection_epoch_atomic = client.connection_epoch_atomic();
518
519        // Queue SetClient (and the instrument cache replay) onto the new
520        // command channel BEFORE publishing it to clones or marking the
521        // connection active. Otherwise a clone observing `is_active()` could
522        // race in and send a Subscribe before SetClient lands, and the
523        // handler would drop the subscription because `inner == None`.
524        let reconnect_handle = client.reconnect_handle();
525        if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
526            anyhow::bail!("Failed to send SetClient command: {e}");
527        }
528
529        if let Some(control) = &self.socket_control {
530            control.register(move || reconnect_handle.request_reconnect());
531        }
532
533        let initial_instruments: Vec<(i64, InstrumentAny)> = self
534            .instruments
535            .iter()
536            .map(|entry| (*entry.key(), entry.value().clone()))
537            .collect();
538
539        if !initial_instruments.is_empty()
540            && let Err(e) = cmd_tx.send(HandlerCommand::InitializeInstruments(initial_instruments))
541        {
542            log::error!("Failed to send InitializeInstruments: {e}");
543        }
544
545        // Publish the new command channel and connection-mode atomic last.
546        // Any clone-driven subscribe call queued from this point lands
547        // behind SetClient and InitializeInstruments in cmd_rx.
548        *self.cmd_tx.write().await = cmd_tx.clone();
549        self.out_rx = Some(out_rx);
550        self.connection_mode.store(connection_mode_atomic);
551        self.connection_epoch.store(connection_epoch_atomic);
552
553        log::debug!("Lighter WebSocket connected: {}", self.url);
554
555        let signal = Arc::clone(&self.signal);
556        let subscriptions = self.subscriptions.clone();
557        let subscription_args = Arc::clone(&self.subscription_args);
558        let cmd_tx_for_reconnect = cmd_tx.clone();
559        let settlement_currency = self.registry.settlement_currency();
560        let book_snapshot_timeout = self.book_snapshot_timeout;
561
562        if let Err(e) = self.task_handle.spawn(async move {
563            let mut handler = FeedHandler::new_with_settlement_currency(
564                Arc::clone(&signal),
565                cmd_rx,
566                raw_rx,
567                out_tx,
568                subscriptions,
569                settlement_currency,
570                book_snapshot_timeout,
571            );
572
573            handler.set_command_sender(cmd_tx_for_reconnect.clone());
574
575            let restore_subscriptions = || {
576                if subscription_args.is_empty() {
577                    log::debug!("No active Lighter subscriptions to restore after reconnect");
578                    return;
579                }
580                log::debug!(
581                    "Restoring {} Lighter subscriptions after reconnect",
582                    subscription_args.len(),
583                );
584
585                // Replay first; the execution client replaces account tokens after reconnect
586                for entry in subscription_args.iter() {
587                    let args = entry.value().clone();
588                    if let Err(e) = cmd_tx_for_reconnect.send(HandlerCommand::Subscribe {
589                        channel: args.channel,
590                        auth: args.auth,
591                        response_tx: None,
592                    }) {
593                        log::error!("Failed to resend Lighter subscribe command: {e}");
594                    }
595                }
596            };
597
598            loop {
599                match handler.next().await {
600                    Some(NautilusWsMessage::Reconnected { connection_epoch }) => {
601                        log::debug!("Lighter WebSocket reconnected");
602                        restore_subscriptions();
603
604                        if handler
605                            .send(NautilusWsMessage::Reconnected { connection_epoch })
606                            .is_err()
607                        {
608                            if handler.is_stopped() {
609                                log::debug!("Failed to forward Reconnected (receiver dropped)");
610                            } else {
611                                log::error!("Failed to forward Reconnected (receiver dropped)");
612                            }
613                            break;
614                        }
615                    }
616                    Some(msg) => {
617                        if handler.send(msg).is_err() {
618                            if handler.is_stopped() {
619                                log::debug!("Failed to send Lighter message (receiver dropped)");
620                            } else {
621                                log::error!("Failed to send Lighter message (receiver dropped)");
622                            }
623                            break;
624                        }
625                    }
626                    None => {
627                        if handler.is_stopped() {
628                            log::debug!("Lighter handler stop signal observed, exiting loop");
629                            break;
630                        }
631                        log::warn!("Lighter WebSocket stream ended unexpectedly");
632                        break;
633                    }
634                }
635            }
636            log::debug!("Lighter handler task completed");
637        }) {
638            self.out_rx = None;
639            anyhow::bail!("Failed to start Lighter WebSocket handler task: {e}");
640        }
641        Ok(())
642    }
643
644    /// Disconnects gracefully: signals shutdown, drains the handler, then
645    /// awaits the task handle with a timeout.
646    ///
647    /// # Errors
648    ///
649    /// This function currently completes best-effort shutdown and returns `Ok(())`.
650    pub async fn disconnect(&mut self) -> Result<(), LighterWsError> {
651        self.connection_generation.fetch_add(1, Ordering::AcqRel);
652        self.initial_connect_cancellation.load().cancel();
653
654        let _guard = self.connection_lock.lock().await;
655        self.initial_connect_cancellation.load().cancel();
656
657        log::debug!("Disconnecting Lighter WebSocket");
658
659        if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
660            log::debug!("Failed to send Lighter disconnect command: {e}");
661        }
662        self.signal.store(true, Ordering::Release);
663
664        let task_result = match finish_task(
665            &mut self.task_handle,
666            DISCONNECT_TIMEOUT,
667            DISCONNECT_TIMEOUT,
668        )
669        .await
670        {
671            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => Ok(()),
672            Some(TaskJoinOutcome::Failed(error)) => Err(LighterWsError::Client(format!(
673                "WebSocket handler task failed: {error}"
674            ))),
675            Some(TaskJoinOutcome::Incomplete) => Err(LighterWsError::Client(
676                "WebSocket handler task did not stop after abort".to_string(),
677            )),
678        };
679
680        self.connection_mode
681            .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));
682
683        if let Some(control) = &self.socket_control {
684            control.deregister();
685        }
686        task_result
687    }
688
689    pub(crate) async fn disconnect_with_task_retention(
690        self,
691        retained: Arc<RetainedTaskSlot>,
692    ) -> Result<(), LighterWsError> {
693        let mut guard = TaskRetentionGuard::new(self, retained);
694
695        guard.client_mut().disconnect().await
696    }
697
698    pub(crate) fn begin_shutdown(&self) {
699        self.initial_connect_cancellation.load().cancel();
700        self.signal.store(true, Ordering::Release);
701    }
702
703    /// Receives the next message from the handler, or `None` if the receiver
704    /// has been taken or the handler has shut down.
705    pub async fn next_event(&mut self) -> Option<NautilusWsMessage> {
706        if let Some(rx) = self.out_rx.as_mut() {
707            rx.recv().await
708        } else {
709            None
710        }
711    }
712
713    /// Takes the feed-handler task slot, leaving an empty slot behind.
714    ///
715    /// Used by callers that connect on a cloned client and want to await the
716    /// inner handler task on a different instance during disconnect.
717    #[must_use]
718    pub(crate) fn take_task_slot(&mut self) -> TaskSlot<()> {
719        std::mem::take(&mut self.task_handle)
720    }
721
722    /// Installs a feed-handler task slot previously obtained from [`Self::take_task_slot`].
723    pub(crate) fn set_task_slot(&mut self, slot: TaskSlot<()>) {
724        assert!(self.task_handle.is_none(), "task slot is already occupied");
725        self.task_handle = slot;
726    }
727
728    /// Subscribe to L2 order-book updates for an instrument.
729    ///
730    /// # Errors
731    ///
732    /// Returns an error if the instrument is not registered, the command
733    /// cannot be queued, or the venue rejects the subscription.
734    pub async fn subscribe_book(&self, instrument_id: InstrumentId) -> Result<(), LighterWsError> {
735        let market_index = self.market_index_for(&instrument_id)?;
736        self.send_cmd(HandlerCommand::SetBookDeltasSub {
737            market_index,
738            subscribed: true,
739        })
740        .await?;
741
742        if let Err(e) = self.subscribe_order_book_stream(market_index).await {
743            let _ = self
744                .send_cmd(HandlerCommand::SetBookDeltasSub {
745                    market_index,
746                    subscribed: false,
747                })
748                .await;
749            return Err(e);
750        }
751
752        Ok(())
753    }
754
755    /// Unsubscribe from L2 order-book updates.
756    ///
757    /// # Errors
758    ///
759    /// Returns an error if the instrument is not registered or the command
760    /// cannot be queued.
761    pub async fn unsubscribe_book(
762        &self,
763        instrument_id: InstrumentId,
764    ) -> Result<(), LighterWsError> {
765        let market_index = self.market_index_for(&instrument_id)?;
766        self.send_cmd(HandlerCommand::SetBookDeltasSub {
767            market_index,
768            subscribed: false,
769        })
770        .await?;
771        self.unsubscribe_order_book_stream(market_index).await
772    }
773
774    /// Subscribe to depth-10 snapshots derived from the same `order_book`
775    /// stream as [`Self::subscribe_book`].
776    ///
777    /// # Errors
778    ///
779    /// Returns an error if the instrument is not registered, the command
780    /// cannot be queued, or the venue rejects the subscription.
781    pub async fn subscribe_book_depth(
782        &self,
783        instrument_id: InstrumentId,
784    ) -> Result<(), LighterWsError> {
785        let market_index = self.market_index_for(&instrument_id)?;
786        self.send_cmd(HandlerCommand::SetDepthSub {
787            market_index,
788            subscribed: true,
789        })
790        .await?;
791
792        if let Err(e) = self.subscribe_order_book_stream(market_index).await {
793            let _ = self
794                .send_cmd(HandlerCommand::SetDepthSub {
795                    market_index,
796                    subscribed: false,
797                })
798                .await;
799            return Err(e);
800        }
801
802        Ok(())
803    }
804
805    /// Unsubscribe from depth-10 snapshots.
806    ///
807    /// Clears the depth-10 emission flag without tearing down the underlying
808    /// `order_book` stream so any active deltas subscriber keeps receiving
809    /// updates.
810    ///
811    /// # Errors
812    ///
813    /// Returns an error if the instrument is not registered or the command
814    /// cannot be queued.
815    pub async fn unsubscribe_book_depth(
816        &self,
817        instrument_id: InstrumentId,
818    ) -> Result<(), LighterWsError> {
819        let market_index = self.market_index_for(&instrument_id)?;
820        self.send_cmd(HandlerCommand::SetDepthSub {
821            market_index,
822            subscribed: false,
823        })
824        .await?;
825        self.unsubscribe_order_book_stream(market_index).await
826    }
827
828    /// Subscribe to ticker (best bid/offer) updates.
829    ///
830    /// # Errors
831    ///
832    /// Returns an error if the instrument is not registered, the command
833    /// cannot be queued, or the venue rejects the subscription.
834    pub async fn subscribe_quotes(
835        &self,
836        instrument_id: InstrumentId,
837    ) -> Result<(), LighterWsError> {
838        let market_index = self.market_index_for(&instrument_id)?;
839        self.send_subscribe(LighterWsChannel::Ticker(market_index), None)
840            .await
841    }
842
843    /// Unsubscribe from ticker updates.
844    ///
845    /// # Errors
846    ///
847    /// Returns an error if the instrument is not registered or the command
848    /// cannot be queued.
849    pub async fn unsubscribe_quotes(
850        &self,
851        instrument_id: InstrumentId,
852    ) -> Result<(), LighterWsError> {
853        let market_index = self.market_index_for(&instrument_id)?;
854        self.send_unsubscribe(LighterWsChannel::Ticker(market_index))
855            .await
856    }
857
858    /// Subscribe to trade updates.
859    ///
860    /// # Errors
861    ///
862    /// Returns an error if the instrument is not registered, the command
863    /// cannot be queued, or the venue rejects the subscription.
864    pub async fn subscribe_trades(
865        &self,
866        instrument_id: InstrumentId,
867    ) -> Result<(), LighterWsError> {
868        let market_index = self.market_index_for(&instrument_id)?;
869        self.send_subscribe(LighterWsChannel::Trade(market_index), None)
870            .await
871    }
872
873    /// Unsubscribe from trade updates.
874    ///
875    /// # Errors
876    ///
877    /// Returns an error if the instrument is not registered or the command
878    /// cannot be queued.
879    pub async fn unsubscribe_trades(
880        &self,
881        instrument_id: InstrumentId,
882    ) -> Result<(), LighterWsError> {
883        let market_index = self.market_index_for(&instrument_id)?;
884        self.send_unsubscribe(LighterWsChannel::Trade(market_index))
885            .await
886    }
887
888    /// Subscribe to the `candle/{market_id}/{resolution}` stream for an
889    /// instrument and resolution.
890    ///
891    /// # Errors
892    ///
893    /// Returns an error if the instrument is not registered, the resolution
894    /// is not offered on the WebSocket stream, the command cannot be queued,
895    /// or the venue rejects the subscription.
896    pub async fn subscribe_candles(
897        &self,
898        instrument_id: InstrumentId,
899        resolution: LighterCandleResolution,
900    ) -> Result<(), LighterWsError> {
901        if !resolution.is_ws_streamable() {
902            return Err(LighterWsError::Client(format!(
903                "resolution {resolution:?} is not offered on the Lighter candle WebSocket stream",
904            )));
905        }
906        let market_index = self.market_index_for(&instrument_id)?;
907        self.send_subscribe(
908            LighterWsChannel::Candle {
909                market_index,
910                resolution,
911            },
912            None,
913        )
914        .await
915    }
916
917    /// Unsubscribe from a candle stream.
918    ///
919    /// # Errors
920    ///
921    /// Returns an error if the instrument is not registered or the command
922    /// cannot be queued.
923    pub async fn unsubscribe_candles(
924        &self,
925        instrument_id: InstrumentId,
926        resolution: LighterCandleResolution,
927    ) -> Result<(), LighterWsError> {
928        let market_index = self.market_index_for(&instrument_id)?;
929        self.send_unsubscribe(LighterWsChannel::Candle {
930            market_index,
931            resolution,
932        })
933        .await
934    }
935
936    /// Subscribe to a market-stats stream covering all markets or a single
937    /// market index.
938    ///
939    /// # Errors
940    ///
941    /// Returns an error if the command cannot be queued or the venue rejects
942    /// the subscription.
943    pub async fn subscribe_market_stats(
944        &self,
945        selection: LighterMarketSelection,
946    ) -> Result<(), LighterWsError> {
947        self.send_subscribe(LighterWsChannel::MarketStats(selection), None)
948            .await
949    }
950
951    /// Unsubscribe from a market-stats stream.
952    ///
953    /// # Errors
954    ///
955    /// Returns an error if the command cannot be queued.
956    pub async fn unsubscribe_market_stats(
957        &self,
958        selection: LighterMarketSelection,
959    ) -> Result<(), LighterWsError> {
960        self.send_unsubscribe(LighterWsChannel::MarketStats(selection))
961            .await
962    }
963
964    /// Subscribe to a spot market-stats stream covering all spot markets or a
965    /// single spot market index.
966    ///
967    /// # Errors
968    ///
969    /// Returns an error if the command cannot be queued or the venue rejects
970    /// the subscription.
971    pub async fn subscribe_spot_market_stats(
972        &self,
973        selection: LighterMarketSelection,
974    ) -> Result<(), LighterWsError> {
975        self.send_subscribe(LighterWsChannel::SpotMarketStats(selection), None)
976            .await
977    }
978
979    /// Unsubscribe from a spot market-stats stream.
980    ///
981    /// # Errors
982    ///
983    /// Returns an error if the command cannot be queued.
984    pub async fn unsubscribe_spot_market_stats(
985        &self,
986        selection: LighterMarketSelection,
987    ) -> Result<(), LighterWsError> {
988        self.send_unsubscribe(LighterWsChannel::SpotMarketStats(selection))
989            .await
990    }
991
992    /// Subscribe to the chain-height stream.
993    ///
994    /// # Errors
995    ///
996    /// Returns an error if the command cannot be queued or the venue rejects
997    /// the subscription.
998    pub async fn subscribe_height(&self) -> Result<(), LighterWsError> {
999        self.send_subscribe(LighterWsChannel::Height, None).await
1000    }
1001
1002    /// Unsubscribe from the chain-height stream.
1003    ///
1004    /// # Errors
1005    ///
1006    /// Returns an error if the command cannot be queued.
1007    pub async fn unsubscribe_height(&self) -> Result<(), LighterWsError> {
1008        self.send_unsubscribe(LighterWsChannel::Height).await
1009    }
1010
1011    /// Provides the execution context the feed handler stamps onto reports
1012    /// parsed from `account_*` frames.
1013    ///
1014    /// Without this context account frames fall back to
1015    /// [`NautilusWsMessage::Raw`]; once it is set the handler emits typed
1016    /// [`crate::websocket::messages::ExecutionReport`] and
1017    /// [`crate::websocket::messages::NautilusWsMessage::AccountState`]
1018    /// messages stamped with `account_id`. The `account_index` is used by
1019    /// the fill parser to determine which side of each account-trade frame
1020    /// the configured account took.
1021    ///
1022    /// # Errors
1023    ///
1024    /// Returns an error if the command cannot be queued.
1025    pub async fn set_execution_context(
1026        &self,
1027        account_id: AccountId,
1028        account_index: i64,
1029    ) -> Result<(), LighterWsError> {
1030        self.send_cmd(HandlerCommand::SetExecutionContext {
1031            account_id,
1032            account_index,
1033        })
1034        .await
1035    }
1036
1037    /// Subscribe to a private account channel using a venue auth token.
1038    ///
1039    /// The auth token must be a valid Lighter L2 auth signature; see the
1040    /// `signing` module for token construction. Re-issuing this with a fresh
1041    /// token (as the execution client's auth-token rotation does) overwrites
1042    /// the stored reconnect-replay token, keeping it within the venue TTL.
1043    ///
1044    /// # Errors
1045    ///
1046    /// Returns an error if the command cannot be queued or the venue rejects
1047    /// the subscription.
1048    pub async fn subscribe_account(
1049        &self,
1050        channel: LighterWsChannel,
1051        auth_token: SecretString,
1052    ) -> Result<(), LighterWsError> {
1053        self.send_subscribe(channel, Some(auth_token)).await
1054    }
1055
1056    /// Unsubscribe from a private account channel.
1057    ///
1058    /// # Errors
1059    ///
1060    /// Returns an error if the command cannot be queued.
1061    pub async fn unsubscribe_account(
1062        &self,
1063        channel: LighterWsChannel,
1064    ) -> Result<(), LighterWsError> {
1065        self.send_unsubscribe(channel).await
1066    }
1067
1068    /// Dispatch a signed L2 transaction over the WebSocket.
1069    ///
1070    /// `tx_type` is the venue's [`crate::common::enums::LighterTxType`]
1071    /// discriminant; `tx_info` is the JSON body produced by the matching
1072    /// [`crate::signing::tx::TxInfoJson`] renderer. The venue confirms
1073    /// acceptance via the `account_*` streams.
1074    ///
1075    /// # Errors
1076    ///
1077    /// Returns an error if the command cannot be queued or the handler cannot
1078    /// report whether it handed the frame to the network writer.
1079    pub async fn send_tx(
1080        &self,
1081        tx_type: u8,
1082        tx_info: Box<serde_json::value::RawValue>,
1083    ) -> Result<(), LighterWsError> {
1084        self.send_tx_on_connection(tx_type, tx_info, self.connection_epoch())
1085            .await
1086    }
1087
1088    pub(crate) async fn send_tx_on_connection(
1089        &self,
1090        tx_type: u8,
1091        tx_info: Box<serde_json::value::RawValue>,
1092        connection_epoch: u64,
1093    ) -> Result<(), LighterWsError> {
1094        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1095        self.send_cmd(HandlerCommand::SendTx {
1096            tx_type,
1097            tx_info,
1098            connection_epoch,
1099            response_tx,
1100        })
1101        .await?;
1102
1103        response_rx.await.map_err(|e| {
1104            LighterWsError::SendTxOutcomeUnknown(format!(
1105                "handler dropped sendTx result after accepting the command: {e}",
1106            ))
1107        })?
1108    }
1109
1110    pub(crate) async fn send_tx_batch_on_connection(
1111        &self,
1112        data: super::messages::LighterWsSendTxBatch,
1113        connection_epoch: u64,
1114    ) -> Result<(), LighterWsError> {
1115        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1116        self.send_cmd(HandlerCommand::SendTxBatch {
1117            data,
1118            connection_epoch,
1119            response_tx,
1120        })
1121        .await?;
1122
1123        response_rx.await.map_err(|e| {
1124            LighterWsError::SendTxOutcomeUnknown(format!(
1125                "handler dropped sendTxBatch result after accepting the command: {e}",
1126            ))
1127        })?
1128    }
1129
1130    async fn send_subscribe(
1131        &self,
1132        channel: LighterWsChannel,
1133        auth: Option<SecretString>,
1134    ) -> Result<(), LighterWsError> {
1135        let topic = channel.topic_key();
1136        let generation = self
1137            .next_subscription_generation
1138            .fetch_add(1, Ordering::Relaxed);
1139        let previous = self.subscription_args.insert(
1140            topic.clone(),
1141            SubscriptionArgs {
1142                channel: channel.clone(),
1143                auth: auth.clone(),
1144                generation,
1145            },
1146        );
1147        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1148
1149        if let Err(e) = self
1150            .send_cmd(HandlerCommand::Subscribe {
1151                channel: channel.clone(),
1152                auth,
1153                response_tx: Some(response_tx),
1154            })
1155            .await
1156        {
1157            if matches!(channel, LighterWsChannel::OrderBook(_)) {
1158                self.remove_subscription_args(&channel, generation);
1159            } else {
1160                self.restore_subscription_args(&topic, generation, previous);
1161            }
1162
1163            return Err(e);
1164        }
1165
1166        match response_rx.await {
1167            Ok(Ok(())) => Ok(()),
1168            Ok(Err(message)) => {
1169                self.remove_subscription_args(&channel, generation);
1170                Err(LighterWsError::Client(message))
1171            }
1172            Err(e) => {
1173                self.remove_subscription_args(&channel, generation);
1174                Err(LighterWsError::Client(format!(
1175                    "handler dropped subscription result for {topic}: {e}",
1176                )))
1177            }
1178        }
1179    }
1180
1181    fn restore_subscription_args(
1182        &self,
1183        topic: &str,
1184        generation: u64,
1185        previous: Option<SubscriptionArgs>,
1186    ) {
1187        let Entry::Occupied(mut entry) = self.subscription_args.entry(topic.to_string()) else {
1188            return;
1189        };
1190
1191        if entry.get().generation != generation {
1192            return;
1193        }
1194
1195        if let Some(previous) = previous {
1196            entry.insert(previous);
1197        } else {
1198            entry.remove();
1199        }
1200    }
1201
1202    fn remove_subscription_args(&self, channel: &LighterWsChannel, generation: u64) {
1203        let topic = channel.topic_key();
1204        // Remove the failed caller before deciding whether another book consumer needs replay.
1205        if matches!(channel, LighterWsChannel::OrderBook(_))
1206            && !self.subscriptions.remove_reference(&topic)
1207        {
1208            return;
1209        }
1210
1211        let Entry::Occupied(entry) = self.subscription_args.entry(topic) else {
1212            return;
1213        };
1214
1215        if entry.get().generation == generation {
1216            entry.remove();
1217        }
1218    }
1219
1220    async fn send_unsubscribe(&self, channel: LighterWsChannel) -> Result<(), LighterWsError> {
1221        let topic = channel.topic_key();
1222        self.send_cmd(HandlerCommand::Unsubscribe { channel })
1223            .await?;
1224        self.subscription_args.remove(&topic);
1225        Ok(())
1226    }
1227
1228    async fn subscribe_order_book_stream(&self, market_index: i64) -> Result<(), LighterWsError> {
1229        let channel = LighterWsChannel::OrderBook(market_index);
1230        let topic = channel.topic_key();
1231
1232        if !self.subscriptions.add_reference(topic.as_str()) {
1233            return Ok(());
1234        }
1235
1236        self.send_subscribe(channel, None).await
1237    }
1238
1239    async fn unsubscribe_order_book_stream(&self, market_index: i64) -> Result<(), LighterWsError> {
1240        let channel = LighterWsChannel::OrderBook(market_index);
1241        let topic = channel.topic_key();
1242
1243        if !self.subscriptions.remove_reference(topic.as_str()) {
1244            return Ok(());
1245        }
1246
1247        if let Err(e) = self.send_unsubscribe(channel).await {
1248            self.subscriptions.add_reference(topic.as_str());
1249            return Err(e);
1250        }
1251
1252        Ok(())
1253    }
1254
1255    async fn send_cmd(&self, cmd: HandlerCommand) -> Result<(), LighterWsError> {
1256        self.cmd_tx
1257            .read()
1258            .await
1259            .send(cmd)
1260            .map_err(|e| LighterWsError::Client(format!("handler unavailable: {e}")))
1261    }
1262
1263    fn market_index_for(&self, instrument_id: &InstrumentId) -> Result<i64, LighterWsError> {
1264        self.registry.market_index(instrument_id).ok_or_else(|| {
1265            LighterWsError::Client(format!(
1266                "no Lighter market_index registered for instrument: {instrument_id}"
1267            ))
1268        })
1269    }
1270}
1271
1272impl Drop for LighterWebSocketClient {
1273    fn drop(&mut self) {
1274        if self.task_handle.is_none() {
1275            return;
1276        }
1277
1278        self.connection_generation.fetch_add(1, Ordering::AcqRel);
1279        self.initial_connect_cancellation.load().cancel();
1280        self.signal.store(true, Ordering::Release);
1281
1282        if let Some(handle) = self.task_handle.as_ref() {
1283            handle.abort();
1284        }
1285
1286        if let Some(control) = &self.socket_control {
1287            control.deregister();
1288        }
1289    }
1290}
1291
1292#[cfg(test)]
1293mod tests {
1294    use nautilus_core::{UnixNanos, string::secret::REDACTED};
1295    use nautilus_model::{
1296        identifiers::Symbol,
1297        instruments::CryptoPerpetual,
1298        types::{Currency, Price, Quantity},
1299    };
1300    use rstest::rstest;
1301
1302    use super::*;
1303    use crate::common::{
1304        consts::LIGHTER_VENUE,
1305        enums::{LighterProductType, LighterTxType},
1306    };
1307
1308    impl LighterWebSocketClient {
1309        pub(crate) async fn drop_next_send_tx_result_for_test(&self) {
1310            let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1311            *self.cmd_tx.write().await = cmd_tx;
1312
1313            get_runtime().spawn(async move {
1314                if let Some(
1315                    HandlerCommand::SendTx { response_tx, .. }
1316                    | HandlerCommand::SendTxBatch { response_tx, .. },
1317                ) = cmd_rx.recv().await
1318                {
1319                    drop(response_tx);
1320                }
1321            });
1322        }
1323    }
1324
1325    fn registry_with(
1326        market_index: i64,
1327        symbol: &str,
1328        product: LighterProductType,
1329    ) -> Arc<MarketRegistry> {
1330        let registry = Arc::new(MarketRegistry::new());
1331        registry.insert(market_index, symbol, product);
1332        registry
1333    }
1334
1335    #[rstest]
1336    #[case::only_caller(1, false, false)]
1337    #[case::other_consumer(2, false, true)]
1338    #[case::other_consumer_leaves(2, true, false)]
1339    fn book_failed_initial_waiter_preserves_other_consumer_replay(
1340        #[case] references: usize,
1341        #[case] other_leaves: bool,
1342        #[case] retained: bool,
1343    ) {
1344        let client = LighterWebSocketClient::new(
1345            Some("wss://example/test".into()),
1346            LighterEnvironment::Testnet,
1347            registry_with(0, "ETH", LighterProductType::Perp),
1348            TransportBackend::default(),
1349            30,
1350            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1351            None,
1352        );
1353        let topic = "order_book:0";
1354        for _ in 0..references {
1355            client.subscriptions.add_reference(topic);
1356        }
1357
1358        client.subscription_args.insert(
1359            topic.into(),
1360            SubscriptionArgs {
1361                channel: LighterWsChannel::OrderBook(0),
1362                auth: None,
1363                generation: 7,
1364            },
1365        );
1366
1367        if other_leaves {
1368            assert!(!client.subscriptions.remove_reference(topic));
1369        }
1370
1371        client.remove_subscription_args(&LighterWsChannel::OrderBook(0), 7);
1372
1373        assert_eq!(client.subscription_args.contains_key(topic), retained);
1374        assert_eq!(
1375            client.subscriptions.get_reference_count(topic),
1376            references - 1 - usize::from(other_leaves)
1377        );
1378
1379        if retained {
1380            assert_eq!(client.subscription_args.get(topic).unwrap().generation, 7);
1381        }
1382    }
1383
1384    #[rstest]
1385    fn debug_redacts_retained_auth_and_proxy_url() {
1386        let client = LighterWebSocketClient::new(
1387            Some("wss://example/test".to_string()),
1388            LighterEnvironment::Testnet,
1389            Arc::new(MarketRegistry::new()),
1390            TransportBackend::default(),
1391            30,
1392            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1393            Some("http://user:proxy-secret@localhost".to_string()),
1394        );
1395        client.subscription_args.insert(
1396            "account_all:7".to_string(),
1397            SubscriptionArgs {
1398                channel: LighterWsChannel::AccountAll(7),
1399                auth: Some(SecretString::from("auth-token")),
1400                generation: 1,
1401            },
1402        );
1403
1404        let debug = format!("{client:?}");
1405
1406        assert!(debug.contains(REDACTED));
1407        assert!(!debug.contains("proxy-secret"));
1408        assert!(!debug.contains("auth-token"));
1409    }
1410
1411    #[rstest]
1412    fn market_index_for_returns_registered_index() {
1413        let registry = registry_with(7, "ETH", LighterProductType::Perp);
1414        let client = LighterWebSocketClient::new(
1415            Some("wss://example/test".to_string()),
1416            LighterEnvironment::Testnet,
1417            Arc::clone(&registry),
1418            TransportBackend::default(),
1419            30,
1420            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1421            None,
1422        );
1423        let id = registry.instrument_id(7).expect("registered");
1424        assert_eq!(client.market_index_for(&id).unwrap(), 7);
1425    }
1426
1427    #[rstest]
1428    fn market_index_for_unregistered_returns_error() {
1429        let registry = Arc::new(MarketRegistry::new());
1430        let client = LighterWebSocketClient::new(
1431            Some("wss://example/test".to_string()),
1432            LighterEnvironment::Testnet,
1433            registry,
1434            TransportBackend::default(),
1435            30,
1436            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1437            None,
1438        );
1439        let id = InstrumentId::new(Symbol::from_str_unchecked("UNKNOWN-PERP"), *LIGHTER_VENUE);
1440        assert!(client.market_index_for(&id).is_err());
1441    }
1442
1443    #[rstest]
1444    fn cache_instrument_populates_lookup() {
1445        let registry = registry_with(0, "ETH", LighterProductType::Perp);
1446        let client = LighterWebSocketClient::new(
1447            Some("wss://example/test".to_string()),
1448            LighterEnvironment::Testnet,
1449            Arc::clone(&registry),
1450            TransportBackend::default(),
1451            30,
1452            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1453            None,
1454        );
1455        let id = registry.instrument_id(0).expect("registered");
1456        let instrument = stub_instrument(id);
1457        client.cache_instrument(0, instrument);
1458        assert!(client.instruments_cache().contains_key(&0));
1459    }
1460
1461    #[tokio::test]
1462    async fn wait_until_active_uses_configured_timeout() {
1463        let client = LighterWebSocketClient::new(
1464            Some("wss://example/test".to_string()),
1465            LighterEnvironment::Testnet,
1466            Arc::new(MarketRegistry::new()),
1467            TransportBackend::default(),
1468            0,
1469            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1470            None,
1471        );
1472
1473        let error = client
1474            .wait_until_active()
1475            .await
1476            .expect_err("inactive client should time out");
1477
1478        assert!(error.to_string().contains("timeout after 0 seconds"));
1479    }
1480
1481    #[rstest]
1482    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1483    async fn disconnect_awaits_handler_after_timeout_abort() {
1484        struct NotifyOnDrop {
1485            tx: Option<tokio::sync::oneshot::Sender<()>>,
1486        }
1487
1488        impl Drop for NotifyOnDrop {
1489            fn drop(&mut self) {
1490                if let Some(tx) = self.tx.take() {
1491                    let _ = tx.send(());
1492                }
1493            }
1494        }
1495
1496        let mut client = LighterWebSocketClient::new(
1497            Some("wss://example/test".to_string()),
1498            LighterEnvironment::Testnet,
1499            Arc::new(MarketRegistry::new()),
1500            TransportBackend::default(),
1501            30,
1502            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1503            None,
1504        );
1505        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1506        let (drop_tx, mut drop_rx) = tokio::sync::oneshot::channel();
1507        client.task_handle.insert(get_runtime().spawn(async move {
1508            let _notify = NotifyOnDrop { tx: Some(drop_tx) };
1509            let _ = started_tx.send(());
1510            std::thread::sleep(DISCONNECT_TIMEOUT + Duration::from_millis(250));
1511            std::future::pending::<()>().await;
1512        }));
1513        started_rx.await.expect("handler task started");
1514
1515        client.disconnect().await.expect("disconnect");
1516
1517        assert_eq!(drop_rx.try_recv(), Ok(()));
1518    }
1519
1520    #[tokio::test]
1521    async fn drop_clone_does_not_cancel_handler() {
1522        let mut client = LighterWebSocketClient::new(
1523            Some("wss://example/test".to_string()),
1524            LighterEnvironment::Testnet,
1525            Arc::new(MarketRegistry::new()),
1526            TransportBackend::default(),
1527            30,
1528            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1529            None,
1530        );
1531        client
1532            .task_handle
1533            .insert(get_runtime().spawn(std::future::pending()));
1534        let clone = client.clone();
1535
1536        drop(clone);
1537
1538        assert!(!client.signal.load(Ordering::Acquire));
1539        assert!(
1540            !client
1541                .task_handle
1542                .as_ref()
1543                .expect("handler task")
1544                .is_finished()
1545        );
1546    }
1547
1548    #[tokio::test]
1549    async fn cancelled_moved_disconnect_retains_handler_task() {
1550        let mut client = LighterWebSocketClient::new(
1551            Some("wss://example/test".to_string()),
1552            LighterEnvironment::Testnet,
1553            Arc::new(MarketRegistry::new()),
1554            TransportBackend::default(),
1555            30,
1556            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1557            None,
1558        );
1559        client
1560            .task_handle
1561            .insert(get_runtime().spawn(std::future::pending()));
1562        let signal = Arc::clone(&client.signal);
1563        let retained = Arc::new(RetainedTaskSlot::new());
1564        let disconnect_future = client.disconnect_with_task_retention(Arc::clone(&retained));
1565
1566        let disconnect = get_runtime().spawn(disconnect_future);
1567
1568        while !signal.load(Ordering::Acquire) {
1569            tokio::task::yield_now().await;
1570        }
1571        disconnect.abort();
1572        let _ = disconnect.await;
1573
1574        assert!(!retained.is_empty());
1575        retained.0.abort();
1576        retained.finish().await.expect("retained handler shutdown");
1577        assert!(retained.is_empty());
1578    }
1579
1580    #[tokio::test]
1581    async fn send_tx_reports_unknown_outcome_when_handler_result_is_dropped() {
1582        let client = LighterWebSocketClient::new(
1583            Some("wss://example/test".to_string()),
1584            LighterEnvironment::Testnet,
1585            Arc::new(MarketRegistry::new()),
1586            TransportBackend::default(),
1587            30,
1588            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1589            None,
1590        );
1591        client.drop_next_send_tx_result_for_test().await;
1592        let tx_info = serde_json::value::RawValue::from_string("{}".to_string()).unwrap();
1593
1594        let error = client
1595            .send_tx(LighterTxType::CreateOrder as u8, tx_info)
1596            .await
1597            .expect_err("dropped handler result must be ambiguous");
1598
1599        assert!(matches!(error, LighterWsError::SendTxOutcomeUnknown(_)));
1600    }
1601
1602    #[tokio::test]
1603    async fn subscribe_waits_for_venue_result_and_removes_failed_generation() {
1604        let client = LighterWebSocketClient::new(
1605            Some("wss://example/test".to_string()),
1606            LighterEnvironment::Testnet,
1607            Arc::new(MarketRegistry::new()),
1608            TransportBackend::default(),
1609            30,
1610            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1611            None,
1612        );
1613        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1614        *client.cmd_tx.write().await = cmd_tx;
1615
1616        let subscribe_client = client.clone();
1617        let subscribe = get_runtime().spawn(async move {
1618            subscribe_client
1619                .subscribe_market_stats(LighterMarketSelection::Market(0))
1620                .await
1621        });
1622
1623        let command = cmd_rx.recv().await.expect("subscribe command");
1624        let HandlerCommand::Subscribe {
1625            response_tx: Some(response_tx),
1626            ..
1627        } = command
1628        else {
1629            panic!("expected subscribe command with venue result sender");
1630        };
1631
1632        assert!(!subscribe.is_finished());
1633        assert!(client.subscription_args.contains_key("market_stats:0"));
1634
1635        response_tx
1636            .send(Err("venue rejected subscription".to_string()))
1637            .expect("subscription result receiver");
1638        let error = subscribe
1639            .await
1640            .expect("subscribe task")
1641            .expect_err("failed venue open must fail the caller");
1642
1643        assert!(error.to_string().contains("venue rejected subscription"));
1644        assert!(!client.subscription_args.contains_key("market_stats:0"));
1645    }
1646
1647    #[tokio::test]
1648    async fn failed_older_subscribe_does_not_remove_newer_generation() {
1649        let client = LighterWebSocketClient::new(
1650            Some("wss://example/test".to_string()),
1651            LighterEnvironment::Testnet,
1652            Arc::new(MarketRegistry::new()),
1653            TransportBackend::default(),
1654            30,
1655            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1656            None,
1657        );
1658        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1659        *client.cmd_tx.write().await = cmd_tx;
1660
1661        let older_client = client.clone();
1662        let older = get_runtime().spawn(async move {
1663            older_client
1664                .subscribe_market_stats(LighterMarketSelection::Market(0))
1665                .await
1666        });
1667        let HandlerCommand::Subscribe {
1668            response_tx: Some(older_response),
1669            ..
1670        } = cmd_rx.recv().await.expect("older subscribe command")
1671        else {
1672            panic!("expected older subscribe command with venue result sender");
1673        };
1674
1675        let newer_client = client.clone();
1676        let newer = get_runtime().spawn(async move {
1677            newer_client
1678                .subscribe_market_stats(LighterMarketSelection::Market(0))
1679                .await
1680        });
1681        let HandlerCommand::Subscribe {
1682            response_tx: Some(newer_response),
1683            ..
1684        } = cmd_rx.recv().await.expect("newer subscribe command")
1685        else {
1686            panic!("expected newer subscribe command with venue result sender");
1687        };
1688
1689        older_response
1690            .send(Err("older generation failed".to_string()))
1691            .expect("older result receiver");
1692        older
1693            .await
1694            .expect("older subscribe task")
1695            .expect_err("older generation must fail");
1696        assert!(client.subscription_args.contains_key("market_stats:0"));
1697
1698        newer_response
1699            .send(Err("newer generation failed".to_string()))
1700            .expect("newer result receiver");
1701        newer
1702            .await
1703            .expect("newer subscribe task")
1704            .expect_err("newer generation must fail");
1705        assert!(!client.subscription_args.contains_key("market_stats:0"));
1706    }
1707
1708    fn stub_instrument(id: InstrumentId) -> InstrumentAny {
1709        InstrumentAny::CryptoPerpetual(
1710            CryptoPerpetual::builder()
1711                .instrument_id(id)
1712                .raw_symbol(id.symbol)
1713                .base_currency(Currency::from("ETH"))
1714                .quote_currency(Currency::from("USDC"))
1715                .settlement_currency(Currency::from("USDC"))
1716                .is_inverse(false)
1717                .price_precision(2)
1718                .size_precision(4)
1719                .price_increment(Price::from("0.01"))
1720                .size_increment(Quantity::from("0.0001"))
1721                .ts_event(UnixNanos::default())
1722                .ts_init(UnixNanos::default())
1723                .build()
1724                .unwrap(),
1725        )
1726    }
1727}
nautilus_lighter/websocket/client.rs

nautilus_lighter/websocket/
client.rs
