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    ConnectionUrl, 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: ConnectionUrl,
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: url.into(),
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.expose()
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 an error if the connection does not reach the active state within the configured
349    /// 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    /// A fresh connection clears held subscriptions; callers subscribe again for the new session.
412    /// Transport reconnects, including `SocketControl` requests, preserve and replay them.
413    ///
414    /// # Errors
415    ///
416    /// Returns an error if the connection fails permanently, exhausts its timeout, is cancelled,
417    /// or the handler cannot be initialized.
418    pub async fn connect(&mut self) -> anyhow::Result<()> {
419        self.connect_with_cancellation(CancellationToken::new())
420            .await
421    }
422
423    pub(crate) async fn connect_with_cancellation(
424        &mut self,
425        cancellation_token: CancellationToken,
426    ) -> anyhow::Result<()> {
427        let generation = self.connection_generation.load(Ordering::Acquire);
428        let _guard = self.connection_lock.lock().await;
429
430        anyhow::ensure!(
431            generation == self.connection_generation.load(Ordering::Acquire),
432            "Lighter WebSocket initial connection cancelled",
433        );
434
435        if self.is_active() {
436            log::warn!("Lighter WebSocket already connected");
437            return Ok(());
438        }
439
440        if let Some(outcome) = finish_task(
441            &mut self.task_handle,
442            DISCONNECT_TIMEOUT,
443            DISCONNECT_TIMEOUT,
444        )
445        .await
446        {
447            match outcome {
448                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
449                TaskJoinOutcome::Failed(error) => {
450                    anyhow::bail!("Lighter WebSocket handler failed: {error}");
451                }
452                TaskJoinOutcome::Incomplete => {
453                    anyhow::bail!("Lighter WebSocket handler did not stop after abort");
454                }
455            }
456        }
457
458        self.signal.store(false, Ordering::Release);
459        self.initial_connect_cancellation
460            .store(Arc::new(cancellation_token.clone()));
461
462        let (message_handler, raw_rx) = channel_epoch_message_handler();
463        let headers = create_standard_nautilus_headers();
464
465        let cfg = WebSocketConfig {
466            url: self.url.expose().to_owned(),
467            headers,
468            heartbeat_interval_secs: Some(HEARTBEAT_INTERVAL.as_secs()),
469            heartbeat_payload: None,
470            connect_timeout_ms: Some(self.ws_timeout_secs.saturating_mul(1_000).max(1)),
471            reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
472            reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
473            reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
474            reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
475            reconnect_max_attempts: None,
476            heartbeat_timeout_secs: Some(HEARTBEAT_TIMEOUT.as_secs()),
477            idle_timeout_ms: None,
478            writer_capacity: None,
479            backend: self.transport_backend,
480            proxy_url: self
481                .proxy_url
482                .as_ref()
483                .map(|value| value.expose_secret().to_owned()),
484            max_message_size_bytes: None,
485            max_frame_size_bytes: None,
486        };
487
488        let connect = WebSocketClient::epoch_builder()
489            .config(cfg)
490            .epoch_handler(message_handler)
491            .rate_limiter(ws_message_rate_limiter(self.url.expose()))
492            .initial_connect_retry_policy(Self::initial_connect_retry_policy())
493            .cancellation_token(cancellation_token.clone())
494            .maybe_state_sink(
495                self.socket_control
496                    .as_ref()
497                    .map(SocketControl::sink)
498                    .or_else(|| self.socket_sink.clone()),
499            )
500            .connect();
501        let client =
502            match tokio::time::timeout(Duration::from_secs(self.ws_timeout_secs), connect).await {
503                Ok(result) => result?,
504                Err(_) => anyhow::bail!(
505                    "Lighter WebSocket initial connection timeout after {} seconds",
506                    self.ws_timeout_secs,
507                ),
508            };
509
510        if cancellation_token.is_cancelled()
511            || generation != self.connection_generation.load(Ordering::Acquire)
512        {
513            client.disconnect().await;
514
515            anyhow::bail!("Lighter WebSocket initial connection cancelled");
516        }
517
518        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
519        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<NautilusWsMessage>();
520
521        // Capture the connection-mode atomic before moving `client` into the
522        // SetClient command below.
523        let connection_mode_atomic = client.connection_mode_atomic();
524        let connection_epoch_atomic = client.connection_epoch_atomic();
525
526        // Queue SetClient (and the instrument cache replay) onto the new
527        // command channel BEFORE publishing it to clones or marking the
528        // connection active. Otherwise a clone observing `is_active()` could
529        // race in and send a Subscribe before SetClient lands, and the
530        // handler would drop the subscription because `inner == None`.
531        let reconnect_handle = client.reconnect_handle();
532        if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
533            anyhow::bail!("Failed to send SetClient command: {e}");
534        }
535
536        if let Some(control) = &self.socket_control {
537            control.register(move || reconnect_handle.request_reconnect());
538        }
539
540        let initial_instruments: Vec<(i64, InstrumentAny)> = self
541            .instruments
542            .iter()
543            .map(|entry| (*entry.key(), entry.value().clone()))
544            .collect();
545
546        if !initial_instruments.is_empty()
547            && let Err(e) = cmd_tx.send(HandlerCommand::InitializeInstruments(initial_instruments))
548        {
549            log::error!("Failed to send InitializeInstruments: {e}");
550        }
551
552        // Publish the new command channel and connection-mode atomic last.
553        // Any clone-driven subscribe call queued from this point lands
554        // behind SetClient and InitializeInstruments in cmd_rx.
555        {
556            let mut sender = self.cmd_tx.write().await;
557            self.subscriptions.clear();
558            self.subscription_args.clear();
559            *sender = cmd_tx.clone();
560        }
561
562        self.out_rx = Some(out_rx);
563        self.connection_mode.store(connection_mode_atomic);
564        self.connection_epoch.store(connection_epoch_atomic);
565
566        log::debug!("Lighter WebSocket connected: {}", self.url.origin());
567
568        let signal = Arc::clone(&self.signal);
569        let subscriptions = self.subscriptions.clone();
570        let subscription_args = Arc::clone(&self.subscription_args);
571        let cmd_tx_for_reconnect = cmd_tx.clone();
572        let settlement_currency = self.registry.settlement_currency();
573        let book_snapshot_timeout = self.book_snapshot_timeout;
574
575        if let Err(e) = self.task_handle.spawn(async move {
576            let mut handler = FeedHandler::new_with_settlement_currency(
577                Arc::clone(&signal),
578                cmd_rx,
579                raw_rx,
580                out_tx,
581                subscriptions,
582                settlement_currency,
583                book_snapshot_timeout,
584            );
585
586            handler.set_command_sender(cmd_tx_for_reconnect.clone());
587
588            let restore_subscriptions = || {
589                if subscription_args.is_empty() {
590                    log::debug!("No active Lighter subscriptions to restore after reconnect");
591                    return;
592                }
593                log::debug!(
594                    "Restoring {} Lighter subscriptions after reconnect",
595                    subscription_args.len(),
596                );
597
598                // Replay first; the execution client replaces account tokens after reconnect
599                for entry in subscription_args.iter() {
600                    let args = entry.value().clone();
601                    if let Err(e) = cmd_tx_for_reconnect.send(HandlerCommand::Subscribe {
602                        channel: args.channel,
603                        auth: args.auth,
604                        response_tx: None,
605                    }) {
606                        log::error!("Failed to resend Lighter subscribe command: {e}");
607                    }
608                }
609            };
610
611            loop {
612                match handler.next().await {
613                    Some(NautilusWsMessage::Reconnected { connection_epoch }) => {
614                        log::debug!("Lighter WebSocket reconnected");
615                        restore_subscriptions();
616
617                        if handler
618                            .send(NautilusWsMessage::Reconnected { connection_epoch })
619                            .is_err()
620                        {
621                            if handler.is_stopped() {
622                                log::debug!("Failed to forward Reconnected (receiver dropped)");
623                            } else {
624                                log::error!("Failed to forward Reconnected (receiver dropped)");
625                            }
626                            break;
627                        }
628                    }
629                    Some(msg) => {
630                        if handler.send(msg).is_err() {
631                            if handler.is_stopped() {
632                                log::debug!("Failed to send Lighter message (receiver dropped)");
633                            } else {
634                                log::error!("Failed to send Lighter message (receiver dropped)");
635                            }
636                            break;
637                        }
638                    }
639                    None => {
640                        if handler.is_stopped() {
641                            log::debug!("Lighter handler stop signal observed, exiting loop");
642                            break;
643                        }
644                        log::warn!("Lighter WebSocket stream ended unexpectedly");
645                        break;
646                    }
647                }
648            }
649            log::debug!("Lighter handler task completed");
650        }) {
651            self.out_rx = None;
652            anyhow::bail!("Failed to start Lighter WebSocket handler task: {e}");
653        }
654        Ok(())
655    }
656
657    /// Disconnects gracefully: signals shutdown, drains the handler, then
658    /// awaits the task handle with a timeout.
659    ///
660    /// # Errors
661    ///
662    /// Returns an error if the handler task fails or does not stop after being aborted. Shutdown
663    /// still clears the connection mode and deregisters the socket control.
664    pub async fn disconnect(&mut self) -> Result<(), LighterWsError> {
665        self.connection_generation.fetch_add(1, Ordering::AcqRel);
666        self.initial_connect_cancellation.load().cancel();
667
668        let _guard = self.connection_lock.lock().await;
669        self.initial_connect_cancellation.load().cancel();
670
671        log::debug!("Disconnecting Lighter WebSocket");
672
673        if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
674            log::debug!("Failed to send Lighter disconnect command: {e}");
675        }
676        self.signal.store(true, Ordering::Release);
677
678        let task_result = match finish_task(
679            &mut self.task_handle,
680            DISCONNECT_TIMEOUT,
681            DISCONNECT_TIMEOUT,
682        )
683        .await
684        {
685            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => Ok(()),
686            Some(TaskJoinOutcome::Failed(error)) => Err(LighterWsError::Client(format!(
687                "WebSocket handler task failed: {error}"
688            ))),
689            Some(TaskJoinOutcome::Incomplete) => Err(LighterWsError::Client(
690                "WebSocket handler task did not stop after abort".to_string(),
691            )),
692        };
693
694        self.connection_mode
695            .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));
696
697        if let Some(control) = &self.socket_control {
698            control.deregister();
699        }
700        task_result
701    }
702
703    pub(crate) async fn disconnect_with_task_retention(
704        self,
705        retained: Arc<RetainedTaskSlot>,
706    ) -> Result<(), LighterWsError> {
707        let mut guard = TaskRetentionGuard::new(self, retained);
708
709        guard.client_mut().disconnect().await
710    }
711
712    pub(crate) fn begin_shutdown(&self) {
713        self.initial_connect_cancellation.load().cancel();
714        self.signal.store(true, Ordering::Release);
715    }
716
717    /// Receives the next message from the handler, or `None` if the receiver
718    /// has been taken or the handler has shut down.
719    pub async fn next_event(&mut self) -> Option<NautilusWsMessage> {
720        if let Some(rx) = self.out_rx.as_mut() {
721            rx.recv().await
722        } else {
723            None
724        }
725    }
726
727    /// Takes the feed-handler task slot, leaving an empty slot behind.
728    ///
729    /// Used by callers that connect on a cloned client and want to await the
730    /// inner handler task on a different instance during disconnect.
731    #[must_use]
732    pub(crate) fn take_task_slot(&mut self) -> TaskSlot<()> {
733        std::mem::take(&mut self.task_handle)
734    }
735
736    /// Installs a feed-handler task slot previously obtained from [`Self::take_task_slot`].
737    pub(crate) fn set_task_slot(&mut self, slot: TaskSlot<()>) {
738        assert!(self.task_handle.is_none(), "task slot is already occupied");
739        self.task_handle = slot;
740    }
741
742    /// Subscribe to L2 order-book updates for an instrument.
743    ///
744    /// # Errors
745    ///
746    /// Returns an error if the instrument is not registered, the command
747    /// cannot be queued, or the venue rejects the subscription.
748    pub async fn subscribe_book(&self, instrument_id: InstrumentId) -> Result<(), LighterWsError> {
749        let market_index = self.market_index_for(&instrument_id)?;
750
751        let session = self
752            .send_cmd_session(HandlerCommand::SetBookDeltasSub {
753                market_index,
754                subscribed: true,
755            })
756            .await?;
757
758        if let Err(e) = self
759            .subscribe_order_book_stream(market_index, &session)
760            .await
761        {
762            let _ = self
763                .send_cmd_on_session(
764                    HandlerCommand::SetBookDeltasSub {
765                        market_index,
766                        subscribed: false,
767                    },
768                    &session,
769                )
770                .await;
771            return Err(e);
772        }
773
774        Ok(())
775    }
776
777    /// Unsubscribe from L2 order-book updates.
778    ///
779    /// # Errors
780    ///
781    /// Returns an error if the instrument is not registered or the command
782    /// cannot be queued.
783    pub async fn unsubscribe_book(
784        &self,
785        instrument_id: InstrumentId,
786    ) -> Result<(), LighterWsError> {
787        let market_index = self.market_index_for(&instrument_id)?;
788        self.send_cmd(HandlerCommand::SetBookDeltasSub {
789            market_index,
790            subscribed: false,
791        })
792        .await?;
793        self.unsubscribe_order_book_stream(market_index).await
794    }
795
796    /// Subscribe to depth-10 snapshots derived from the same `order_book`
797    /// stream as [`Self::subscribe_book`].
798    ///
799    /// # Errors
800    ///
801    /// Returns an error if the instrument is not registered, the command
802    /// cannot be queued, or the venue rejects the subscription.
803    pub async fn subscribe_book_depth(
804        &self,
805        instrument_id: InstrumentId,
806    ) -> Result<(), LighterWsError> {
807        let market_index = self.market_index_for(&instrument_id)?;
808
809        let session = self
810            .send_cmd_session(HandlerCommand::SetDepthSub {
811                market_index,
812                subscribed: true,
813            })
814            .await?;
815
816        if let Err(e) = self
817            .subscribe_order_book_stream(market_index, &session)
818            .await
819        {
820            let _ = self
821                .send_cmd_on_session(
822                    HandlerCommand::SetDepthSub {
823                        market_index,
824                        subscribed: false,
825                    },
826                    &session,
827                )
828                .await;
829            return Err(e);
830        }
831
832        Ok(())
833    }
834
835    /// Unsubscribe from depth-10 snapshots.
836    ///
837    /// Clears the depth-10 emission flag without tearing down the underlying
838    /// `order_book` stream so any active deltas subscriber keeps receiving
839    /// updates.
840    ///
841    /// # Errors
842    ///
843    /// Returns an error if the instrument is not registered or the command
844    /// cannot be queued.
845    pub async fn unsubscribe_book_depth(
846        &self,
847        instrument_id: InstrumentId,
848    ) -> Result<(), LighterWsError> {
849        let market_index = self.market_index_for(&instrument_id)?;
850        self.send_cmd(HandlerCommand::SetDepthSub {
851            market_index,
852            subscribed: false,
853        })
854        .await?;
855        self.unsubscribe_order_book_stream(market_index).await
856    }
857
858    /// Subscribe to ticker (best bid/offer) 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_quotes(
865        &self,
866        instrument_id: InstrumentId,
867    ) -> Result<(), LighterWsError> {
868        let market_index = self.market_index_for(&instrument_id)?;
869        self.send_subscribe(LighterWsChannel::Ticker(market_index), None)
870            .await
871    }
872
873    /// Unsubscribe from ticker 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_quotes(
880        &self,
881        instrument_id: InstrumentId,
882    ) -> Result<(), LighterWsError> {
883        let market_index = self.market_index_for(&instrument_id)?;
884        self.send_unsubscribe(LighterWsChannel::Ticker(market_index))
885            .await
886    }
887
888    /// Subscribe to trade updates.
889    ///
890    /// # Errors
891    ///
892    /// Returns an error if the instrument is not registered, the command
893    /// cannot be queued, or the venue rejects the subscription.
894    pub async fn subscribe_trades(
895        &self,
896        instrument_id: InstrumentId,
897    ) -> Result<(), LighterWsError> {
898        let market_index = self.market_index_for(&instrument_id)?;
899        self.send_subscribe(LighterWsChannel::Trade(market_index), None)
900            .await
901    }
902
903    /// Unsubscribe from trade updates.
904    ///
905    /// # Errors
906    ///
907    /// Returns an error if the instrument is not registered or the command
908    /// cannot be queued.
909    pub async fn unsubscribe_trades(
910        &self,
911        instrument_id: InstrumentId,
912    ) -> Result<(), LighterWsError> {
913        let market_index = self.market_index_for(&instrument_id)?;
914        self.send_unsubscribe(LighterWsChannel::Trade(market_index))
915            .await
916    }
917
918    /// Subscribe to the `candle/{market_id}/{resolution}` stream for an
919    /// instrument and resolution.
920    ///
921    /// # Errors
922    ///
923    /// Returns an error if the instrument is not registered, the resolution
924    /// is not offered on the WebSocket stream, the command cannot be queued,
925    /// or the venue rejects the subscription.
926    pub async fn subscribe_candles(
927        &self,
928        instrument_id: InstrumentId,
929        resolution: LighterCandleResolution,
930    ) -> Result<(), LighterWsError> {
931        if !resolution.is_ws_streamable() {
932            return Err(LighterWsError::Client(format!(
933                "resolution {resolution:?} is not offered on the Lighter candle WebSocket stream",
934            )));
935        }
936        let market_index = self.market_index_for(&instrument_id)?;
937        self.send_subscribe(
938            LighterWsChannel::Candle {
939                market_index,
940                resolution,
941            },
942            None,
943        )
944        .await
945    }
946
947    /// Unsubscribe from a candle stream.
948    ///
949    /// # Errors
950    ///
951    /// Returns an error if the instrument is not registered or the command
952    /// cannot be queued.
953    pub async fn unsubscribe_candles(
954        &self,
955        instrument_id: InstrumentId,
956        resolution: LighterCandleResolution,
957    ) -> Result<(), LighterWsError> {
958        let market_index = self.market_index_for(&instrument_id)?;
959        self.send_unsubscribe(LighterWsChannel::Candle {
960            market_index,
961            resolution,
962        })
963        .await
964    }
965
966    /// Subscribe to a market-stats stream covering all markets or a single
967    /// market index.
968    ///
969    /// # Errors
970    ///
971    /// Returns an error if the command cannot be queued or the venue rejects
972    /// the subscription.
973    pub async fn subscribe_market_stats(
974        &self,
975        selection: LighterMarketSelection,
976    ) -> Result<(), LighterWsError> {
977        self.send_subscribe(LighterWsChannel::MarketStats(selection), None)
978            .await
979    }
980
981    /// Unsubscribe from a market-stats stream.
982    ///
983    /// # Errors
984    ///
985    /// Returns an error if the command cannot be queued.
986    pub async fn unsubscribe_market_stats(
987        &self,
988        selection: LighterMarketSelection,
989    ) -> Result<(), LighterWsError> {
990        self.send_unsubscribe(LighterWsChannel::MarketStats(selection))
991            .await
992    }
993
994    /// Subscribe to a spot market-stats stream covering all spot markets or a
995    /// single spot market index.
996    ///
997    /// # Errors
998    ///
999    /// Returns an error if the command cannot be queued or the venue rejects
1000    /// the subscription.
1001    pub async fn subscribe_spot_market_stats(
1002        &self,
1003        selection: LighterMarketSelection,
1004    ) -> Result<(), LighterWsError> {
1005        self.send_subscribe(LighterWsChannel::SpotMarketStats(selection), None)
1006            .await
1007    }
1008
1009    /// Unsubscribe from a spot market-stats stream.
1010    ///
1011    /// # Errors
1012    ///
1013    /// Returns an error if the command cannot be queued.
1014    pub async fn unsubscribe_spot_market_stats(
1015        &self,
1016        selection: LighterMarketSelection,
1017    ) -> Result<(), LighterWsError> {
1018        self.send_unsubscribe(LighterWsChannel::SpotMarketStats(selection))
1019            .await
1020    }
1021
1022    /// Subscribe to the chain-height stream.
1023    ///
1024    /// # Errors
1025    ///
1026    /// Returns an error if the command cannot be queued or the venue rejects
1027    /// the subscription.
1028    pub async fn subscribe_height(&self) -> Result<(), LighterWsError> {
1029        self.send_subscribe(LighterWsChannel::Height, None).await
1030    }
1031
1032    /// Unsubscribe from the chain-height stream.
1033    ///
1034    /// # Errors
1035    ///
1036    /// Returns an error if the command cannot be queued.
1037    pub async fn unsubscribe_height(&self) -> Result<(), LighterWsError> {
1038        self.send_unsubscribe(LighterWsChannel::Height).await
1039    }
1040
1041    /// Provides the execution context the feed handler stamps onto reports
1042    /// parsed from `account_*` frames.
1043    ///
1044    /// Without this context account frames fall back to
1045    /// [`NautilusWsMessage::Raw`]; once it is set the handler emits typed
1046    /// [`crate::websocket::messages::ExecutionReport`] and
1047    /// [`crate::websocket::messages::NautilusWsMessage::AccountState`]
1048    /// messages stamped with `account_id`. The `account_index` is used by
1049    /// the fill parser to determine which side of each account-trade frame
1050    /// the configured account took.
1051    ///
1052    /// # Errors
1053    ///
1054    /// Returns an error if the command cannot be queued.
1055    pub async fn set_execution_context(
1056        &self,
1057        account_id: AccountId,
1058        account_index: i64,
1059    ) -> Result<(), LighterWsError> {
1060        self.send_cmd(HandlerCommand::SetExecutionContext {
1061            account_id,
1062            account_index,
1063        })
1064        .await
1065    }
1066
1067    /// Subscribe to a private account channel using a venue auth token.
1068    ///
1069    /// The auth token must be a valid Lighter L2 auth signature; see the
1070    /// `signing` module for token construction. Re-issuing this with a fresh
1071    /// token (as the execution client's auth-token rotation does) overwrites
1072    /// the stored reconnect-replay token, keeping it within the venue TTL.
1073    ///
1074    /// # Errors
1075    ///
1076    /// Returns an error if the command cannot be queued or the venue rejects
1077    /// the subscription.
1078    pub async fn subscribe_account(
1079        &self,
1080        channel: LighterWsChannel,
1081        auth_token: SecretString,
1082    ) -> Result<(), LighterWsError> {
1083        self.send_subscribe(channel, Some(auth_token)).await
1084    }
1085
1086    /// Unsubscribe from a private account channel.
1087    ///
1088    /// # Errors
1089    ///
1090    /// Returns an error if the command cannot be queued.
1091    pub async fn unsubscribe_account(
1092        &self,
1093        channel: LighterWsChannel,
1094    ) -> Result<(), LighterWsError> {
1095        self.send_unsubscribe(channel).await
1096    }
1097
1098    /// Dispatch a signed L2 transaction over the WebSocket.
1099    ///
1100    /// `tx_type` is the venue's [`crate::common::enums::LighterTxType`]
1101    /// discriminant; `tx_info` is the JSON body produced by the matching
1102    /// [`crate::signing::tx::TxInfoJson`] renderer. The venue confirms
1103    /// acceptance via the `account_*` streams.
1104    ///
1105    /// # Errors
1106    ///
1107    /// Returns an error if the command cannot be queued or the handler cannot
1108    /// report whether it handed the frame to the network writer.
1109    pub async fn send_tx(
1110        &self,
1111        tx_type: u8,
1112        tx_info: Box<serde_json::value::RawValue>,
1113    ) -> Result<(), LighterWsError> {
1114        self.send_tx_on_connection(tx_type, tx_info, self.connection_epoch())
1115            .await
1116    }
1117
1118    pub(crate) async fn send_tx_on_connection(
1119        &self,
1120        tx_type: u8,
1121        tx_info: Box<serde_json::value::RawValue>,
1122        connection_epoch: u64,
1123    ) -> Result<(), LighterWsError> {
1124        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1125        self.send_cmd(HandlerCommand::SendTx {
1126            tx_type,
1127            tx_info,
1128            connection_epoch,
1129            response_tx,
1130        })
1131        .await?;
1132
1133        response_rx.await.map_err(|e| {
1134            LighterWsError::SendTxOutcomeUnknown(format!(
1135                "handler dropped sendTx result after accepting the command: {e}",
1136            ))
1137        })?
1138    }
1139
1140    pub(crate) async fn send_tx_batch_on_connection(
1141        &self,
1142        data: super::messages::LighterWsSendTxBatch,
1143        connection_epoch: u64,
1144    ) -> Result<(), LighterWsError> {
1145        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1146        self.send_cmd(HandlerCommand::SendTxBatch {
1147            data,
1148            connection_epoch,
1149            response_tx,
1150        })
1151        .await?;
1152
1153        response_rx.await.map_err(|e| {
1154            LighterWsError::SendTxOutcomeUnknown(format!(
1155                "handler dropped sendTxBatch result after accepting the command: {e}",
1156            ))
1157        })?
1158    }
1159
1160    async fn send_subscribe(
1161        &self,
1162        channel: LighterWsChannel,
1163        auth: Option<SecretString>,
1164    ) -> Result<(), LighterWsError> {
1165        self.send_subscribe_on_session(channel, auth, None).await
1166    }
1167
1168    async fn send_subscribe_on_session(
1169        &self,
1170        channel: LighterWsChannel,
1171        auth: Option<SecretString>,
1172        expected_session: Option<&tokio::sync::mpsc::UnboundedSender<HandlerCommand>>,
1173    ) -> Result<(), LighterWsError> {
1174        let topic = channel.topic_key();
1175        let sender = self.cmd_tx.read().await;
1176
1177        if expected_session.is_some_and(|session| !sender.same_channel(session)) {
1178            return Err(LighterWsError::Client(format!(
1179                "WebSocket session changed before subscription for {topic}",
1180            )));
1181        }
1182
1183        if matches!(channel, LighterWsChannel::OrderBook(_))
1184            && !self.subscriptions.add_reference(&topic)
1185        {
1186            return Ok(());
1187        }
1188
1189        let generation = self
1190            .next_subscription_generation
1191            .fetch_add(1, Ordering::Relaxed);
1192        let previous = self.subscription_args.insert(
1193            topic.clone(),
1194            SubscriptionArgs {
1195                channel: channel.clone(),
1196                auth: auth.clone(),
1197                generation,
1198            },
1199        );
1200        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1201
1202        if let Err(e) = send_command(
1203            &sender,
1204            HandlerCommand::Subscribe {
1205                channel: channel.clone(),
1206                auth,
1207                response_tx: Some(response_tx),
1208            },
1209        ) {
1210            if matches!(channel, LighterWsChannel::OrderBook(_)) {
1211                self.remove_subscription_args(&channel, generation);
1212            } else {
1213                self.restore_subscription_args(&topic, generation, previous);
1214            }
1215
1216            return Err(e);
1217        }
1218
1219        let session = sender.clone();
1220        drop(sender);
1221
1222        let result = match response_rx.await {
1223            Ok(result) => result.map_err(LighterWsError::Client),
1224            Err(e) => Err(LighterWsError::Client(format!(
1225                "handler dropped subscription result for {topic}: {e}",
1226            ))),
1227        };
1228
1229        let sender = self.cmd_tx.read().await;
1230
1231        if result.is_err() && sender.same_channel(&session) {
1232            self.remove_subscription_args(&channel, generation);
1233        }
1234
1235        result
1236    }
1237
1238    fn restore_subscription_args(
1239        &self,
1240        topic: &str,
1241        generation: u64,
1242        previous: Option<SubscriptionArgs>,
1243    ) {
1244        let Entry::Occupied(mut entry) = self.subscription_args.entry(topic.to_string()) else {
1245            return;
1246        };
1247
1248        if entry.get().generation != generation {
1249            return;
1250        }
1251
1252        if let Some(previous) = previous {
1253            entry.insert(previous);
1254        } else {
1255            entry.remove();
1256        }
1257    }
1258
1259    fn remove_subscription_args(&self, channel: &LighterWsChannel, generation: u64) {
1260        let topic = channel.topic_key();
1261        // Remove the failed caller before deciding whether another book consumer needs replay.
1262        if matches!(channel, LighterWsChannel::OrderBook(_))
1263            && !self.subscriptions.remove_reference(&topic)
1264        {
1265            return;
1266        }
1267
1268        let Entry::Occupied(entry) = self.subscription_args.entry(topic) else {
1269            return;
1270        };
1271
1272        if entry.get().generation == generation {
1273            entry.remove();
1274        }
1275    }
1276
1277    async fn send_unsubscribe(&self, channel: LighterWsChannel) -> Result<(), LighterWsError> {
1278        let topic = channel.topic_key();
1279        self.send_cmd(HandlerCommand::Unsubscribe { channel })
1280            .await?;
1281        self.subscription_args.remove(&topic);
1282        Ok(())
1283    }
1284
1285    async fn subscribe_order_book_stream(
1286        &self,
1287        market_index: i64,
1288        session: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1289    ) -> Result<(), LighterWsError> {
1290        self.send_subscribe_on_session(
1291            LighterWsChannel::OrderBook(market_index),
1292            None,
1293            Some(session),
1294        )
1295        .await
1296    }
1297
1298    async fn unsubscribe_order_book_stream(&self, market_index: i64) -> Result<(), LighterWsError> {
1299        let channel = LighterWsChannel::OrderBook(market_index);
1300        let topic = channel.topic_key();
1301
1302        if !self.subscriptions.remove_reference(topic.as_str()) {
1303            return Ok(());
1304        }
1305
1306        if let Err(e) = self.send_unsubscribe(channel).await {
1307            self.subscriptions.add_reference(topic.as_str());
1308            return Err(e);
1309        }
1310
1311        Ok(())
1312    }
1313
1314    async fn send_cmd(&self, cmd: HandlerCommand) -> Result<(), LighterWsError> {
1315        let sender = self.cmd_tx.read().await;
1316        send_command(&sender, cmd)
1317    }
1318
1319    async fn send_cmd_session(
1320        &self,
1321        cmd: HandlerCommand,
1322    ) -> Result<tokio::sync::mpsc::UnboundedSender<HandlerCommand>, LighterWsError> {
1323        let sender = self.cmd_tx.read().await;
1324        send_command(&sender, cmd)?;
1325        Ok(sender.clone())
1326    }
1327
1328    async fn send_cmd_on_session(
1329        &self,
1330        cmd: HandlerCommand,
1331        session: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1332    ) -> Result<(), LighterWsError> {
1333        let sender = self.cmd_tx.read().await;
1334        if sender.same_channel(session) {
1335            send_command(&sender, cmd)?;
1336        }
1337
1338        Ok(())
1339    }
1340
1341    fn market_index_for(&self, instrument_id: &InstrumentId) -> Result<i64, LighterWsError> {
1342        self.registry.market_index(instrument_id).ok_or_else(|| {
1343            LighterWsError::Client(format!(
1344                "no Lighter market_index registered for instrument: {instrument_id}"
1345            ))
1346        })
1347    }
1348}
1349
1350impl Drop for LighterWebSocketClient {
1351    fn drop(&mut self) {
1352        if self.task_handle.is_none() {
1353            return;
1354        }
1355
1356        self.connection_generation.fetch_add(1, Ordering::AcqRel);
1357        self.initial_connect_cancellation.load().cancel();
1358        self.signal.store(true, Ordering::Release);
1359
1360        if let Some(handle) = self.task_handle.as_ref() {
1361            handle.abort();
1362        }
1363
1364        if let Some(control) = &self.socket_control {
1365            control.deregister();
1366        }
1367    }
1368}
1369
1370fn send_command(
1371    sender: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1372    cmd: HandlerCommand,
1373) -> Result<(), LighterWsError> {
1374    sender
1375        .send(cmd)
1376        .map_err(|e| LighterWsError::Client(format!("handler unavailable: {e}")))
1377}
1378
1379#[cfg(test)]
1380mod tests {
1381    use nautilus_core::{UnixNanos, string::secret::REDACTED};
1382    use nautilus_model::{
1383        identifiers::Symbol,
1384        instruments::CryptoPerpetual,
1385        types::{Currency, Price, Quantity},
1386    };
1387    use rstest::rstest;
1388
1389    use super::*;
1390    use crate::common::{
1391        consts::LIGHTER_VENUE,
1392        enums::{LighterProductType, LighterTxType},
1393    };
1394
1395    impl LighterWebSocketClient {
1396        pub(crate) async fn drop_next_send_tx_result_for_test(&self) {
1397            let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1398            *self.cmd_tx.write().await = cmd_tx;
1399
1400            get_runtime().spawn(async move {
1401                if let Some(
1402                    HandlerCommand::SendTx { response_tx, .. }
1403                    | HandlerCommand::SendTxBatch { response_tx, .. },
1404                ) = cmd_rx.recv().await
1405                {
1406                    cmd_rx.close();
1407                    drop(response_tx);
1408                }
1409            });
1410        }
1411    }
1412
1413    fn registry_with(
1414        market_index: i64,
1415        symbol: &str,
1416        product: LighterProductType,
1417    ) -> Arc<MarketRegistry> {
1418        let registry = Arc::new(MarketRegistry::new());
1419        registry.insert(market_index, symbol, product);
1420        registry
1421    }
1422
1423    #[rstest]
1424    #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
1425    #[case(
1426        "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
1427        "wss://stream.example.com:9443"
1428    )]
1429    fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
1430        let client = LighterWebSocketClient::new(
1431            Some(url.to_string()),
1432            LighterEnvironment::Testnet,
1433            Arc::new(MarketRegistry::new()),
1434            TransportBackend::default(),
1435            30,
1436            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1437            None,
1438        );
1439        let debug = format!("{client:?}");
1440        assert!(debug.contains(&format!("url: {origin:?}")));
1441        assert!(!debug.contains("SYN-"));
1442        assert_eq!(client.url.expose(), url);
1443    }
1444
1445    #[rstest]
1446    #[case::only_caller(1, false, false)]
1447    #[case::other_consumer(2, false, true)]
1448    #[case::other_consumer_leaves(2, true, false)]
1449    fn book_failed_initial_waiter_preserves_other_consumer_replay(
1450        #[case] references: usize,
1451        #[case] other_leaves: bool,
1452        #[case] retained: bool,
1453    ) {
1454        let client = LighterWebSocketClient::new(
1455            Some("wss://example/test".into()),
1456            LighterEnvironment::Testnet,
1457            registry_with(0, "ETH", LighterProductType::Perp),
1458            TransportBackend::default(),
1459            30,
1460            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1461            None,
1462        );
1463        let topic = "order_book:0";
1464        for _ in 0..references {
1465            client.subscriptions.add_reference(topic);
1466        }
1467
1468        client.subscription_args.insert(
1469            topic.into(),
1470            SubscriptionArgs {
1471                channel: LighterWsChannel::OrderBook(0),
1472                auth: None,
1473                generation: 7,
1474            },
1475        );
1476
1477        if other_leaves {
1478            assert!(!client.subscriptions.remove_reference(topic));
1479        }
1480
1481        client.remove_subscription_args(&LighterWsChannel::OrderBook(0), 7);
1482
1483        assert_eq!(client.subscription_args.contains_key(topic), retained);
1484        assert_eq!(
1485            client.subscriptions.get_reference_count(topic),
1486            references - 1 - usize::from(other_leaves)
1487        );
1488
1489        if retained {
1490            assert_eq!(client.subscription_args.get(topic).unwrap().generation, 7);
1491        }
1492    }
1493
1494    #[rstest]
1495    fn debug_redacts_retained_auth_and_proxy_url() {
1496        let 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            Some("http://user:proxy-secret@localhost".to_string()),
1504        );
1505        client.subscription_args.insert(
1506            "account_all:7".to_string(),
1507            SubscriptionArgs {
1508                channel: LighterWsChannel::AccountAll(7),
1509                auth: Some(SecretString::from("auth-token")),
1510                generation: 1,
1511            },
1512        );
1513
1514        let debug = format!("{client:?}");
1515
1516        assert!(debug.contains(REDACTED));
1517        assert!(!debug.contains("proxy-secret"));
1518        assert!(!debug.contains("auth-token"));
1519    }
1520
1521    #[rstest]
1522    fn market_index_for_returns_registered_index() {
1523        let registry = registry_with(7, "ETH", LighterProductType::Perp);
1524        let client = LighterWebSocketClient::new(
1525            Some("wss://example/test".to_string()),
1526            LighterEnvironment::Testnet,
1527            Arc::clone(&registry),
1528            TransportBackend::default(),
1529            30,
1530            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1531            None,
1532        );
1533        let id = registry.instrument_id(7).expect("registered");
1534        assert_eq!(client.market_index_for(&id).unwrap(), 7);
1535    }
1536
1537    #[rstest]
1538    fn market_index_for_unregistered_returns_error() {
1539        let registry = Arc::new(MarketRegistry::new());
1540        let client = LighterWebSocketClient::new(
1541            Some("wss://example/test".to_string()),
1542            LighterEnvironment::Testnet,
1543            registry,
1544            TransportBackend::default(),
1545            30,
1546            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1547            None,
1548        );
1549        let id = InstrumentId::new(Symbol::from_str_unchecked("UNKNOWN-PERP"), *LIGHTER_VENUE);
1550        assert!(client.market_index_for(&id).is_err());
1551    }
1552
1553    #[rstest]
1554    fn cache_instrument_populates_lookup() {
1555        let registry = registry_with(0, "ETH", LighterProductType::Perp);
1556        let client = LighterWebSocketClient::new(
1557            Some("wss://example/test".to_string()),
1558            LighterEnvironment::Testnet,
1559            Arc::clone(&registry),
1560            TransportBackend::default(),
1561            30,
1562            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1563            None,
1564        );
1565        let id = registry.instrument_id(0).expect("registered");
1566        let instrument = stub_instrument(id);
1567        client.cache_instrument(0, instrument);
1568        assert!(client.instruments_cache().contains_key(&0));
1569    }
1570
1571    #[tokio::test]
1572    async fn wait_until_active_uses_configured_timeout() {
1573        let client = LighterWebSocketClient::new(
1574            Some("wss://example/test".to_string()),
1575            LighterEnvironment::Testnet,
1576            Arc::new(MarketRegistry::new()),
1577            TransportBackend::default(),
1578            0,
1579            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1580            None,
1581        );
1582
1583        let error = client
1584            .wait_until_active()
1585            .await
1586            .expect_err("inactive client should time out");
1587
1588        assert!(error.to_string().contains("timeout after 0 seconds"));
1589    }
1590
1591    #[rstest]
1592    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
1593    async fn disconnect_awaits_handler_after_timeout_abort() {
1594        struct NotifyOnDrop {
1595            tx: Option<tokio::sync::oneshot::Sender<()>>,
1596        }
1597
1598        impl Drop for NotifyOnDrop {
1599            fn drop(&mut self) {
1600                if let Some(tx) = self.tx.take() {
1601                    let _ = tx.send(());
1602                }
1603            }
1604        }
1605
1606        let mut client = LighterWebSocketClient::new(
1607            Some("wss://example/test".to_string()),
1608            LighterEnvironment::Testnet,
1609            Arc::new(MarketRegistry::new()),
1610            TransportBackend::default(),
1611            30,
1612            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1613            None,
1614        );
1615        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
1616        let (drop_tx, mut drop_rx) = tokio::sync::oneshot::channel();
1617        client.task_handle.insert(get_runtime().spawn(async move {
1618            let _notify = NotifyOnDrop { tx: Some(drop_tx) };
1619            let _ = started_tx.send(());
1620            std::thread::sleep(DISCONNECT_TIMEOUT + Duration::from_millis(250));
1621            std::future::pending::<()>().await;
1622        }));
1623        started_rx.await.expect("handler task started");
1624
1625        client.disconnect().await.expect("disconnect");
1626
1627        assert_eq!(drop_rx.try_recv(), Ok(()));
1628    }
1629
1630    #[tokio::test]
1631    async fn drop_clone_does_not_cancel_handler() {
1632        let mut client = LighterWebSocketClient::new(
1633            Some("wss://example/test".to_string()),
1634            LighterEnvironment::Testnet,
1635            Arc::new(MarketRegistry::new()),
1636            TransportBackend::default(),
1637            30,
1638            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1639            None,
1640        );
1641        client
1642            .task_handle
1643            .insert(get_runtime().spawn(std::future::pending()));
1644        let clone = client.clone();
1645
1646        drop(clone);
1647
1648        assert!(!client.signal.load(Ordering::Acquire));
1649        assert!(
1650            !client
1651                .task_handle
1652                .as_ref()
1653                .expect("handler task")
1654                .is_finished()
1655        );
1656    }
1657
1658    #[tokio::test]
1659    async fn cancelled_moved_disconnect_retains_handler_task() {
1660        let mut client = LighterWebSocketClient::new(
1661            Some("wss://example/test".to_string()),
1662            LighterEnvironment::Testnet,
1663            Arc::new(MarketRegistry::new()),
1664            TransportBackend::default(),
1665            30,
1666            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1667            None,
1668        );
1669        client
1670            .task_handle
1671            .insert(get_runtime().spawn(std::future::pending()));
1672        let signal = Arc::clone(&client.signal);
1673        let retained = Arc::new(RetainedTaskSlot::new());
1674        let disconnect_future = client.disconnect_with_task_retention(Arc::clone(&retained));
1675
1676        let disconnect = get_runtime().spawn(disconnect_future);
1677
1678        while !signal.load(Ordering::Acquire) {
1679            tokio::task::yield_now().await;
1680        }
1681        disconnect.abort();
1682        let _ = disconnect.await;
1683
1684        assert!(!retained.is_empty());
1685        retained.0.abort();
1686        retained.finish().await.expect("retained handler shutdown");
1687        assert!(retained.is_empty());
1688    }
1689
1690    #[tokio::test]
1691    async fn send_tx_reports_unknown_outcome_when_handler_result_is_dropped() {
1692        let client = LighterWebSocketClient::new(
1693            Some("wss://example/test".to_string()),
1694            LighterEnvironment::Testnet,
1695            Arc::new(MarketRegistry::new()),
1696            TransportBackend::default(),
1697            30,
1698            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1699            None,
1700        );
1701        client.drop_next_send_tx_result_for_test().await;
1702        let tx_info = serde_json::value::RawValue::from_string("{}".to_string()).unwrap();
1703
1704        let error = client
1705            .send_tx(LighterTxType::CreateOrder as u8, tx_info)
1706            .await
1707            .expect_err("dropped handler result must be ambiguous");
1708
1709        assert!(matches!(error, LighterWsError::SendTxOutcomeUnknown(_)));
1710    }
1711
1712    #[tokio::test]
1713    async fn subscription_registration_waits_for_session_publication() {
1714        let client = LighterWebSocketClient::new(
1715            Some("wss://example/test".to_string()),
1716            LighterEnvironment::Testnet,
1717            Arc::new(MarketRegistry::new()),
1718            TransportBackend::default(),
1719            30,
1720            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1721            None,
1722        );
1723        let mut publication = client.cmd_tx.write().await;
1724        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1725        let mut subscribe =
1726            Box::pin(client.subscribe_market_stats(LighterMarketSelection::Market(0)));
1727
1728        assert!(futures_util::poll!(subscribe.as_mut()).is_pending());
1729        assert!(client.subscription_args.is_empty());
1730
1731        client.subscriptions.clear();
1732        client.subscription_args.clear();
1733        *publication = cmd_tx;
1734        drop(publication);
1735
1736        assert!(futures_util::poll!(subscribe.as_mut()).is_pending());
1737        let command = cmd_rx.try_recv().expect("new-session subscribe command");
1738
1739        let HandlerCommand::Subscribe {
1740            channel,
1741            auth,
1742            response_tx: Some(response_tx),
1743        } = command
1744        else {
1745            panic!("expected subscribe command with venue result sender");
1746        };
1747
1748        response_tx
1749            .send(Ok(()))
1750            .expect("subscription result receiver");
1751        subscribe.await.expect("new-session subscription");
1752
1753        assert_eq!(
1754            channel,
1755            LighterWsChannel::MarketStats(LighterMarketSelection::Market(0))
1756        );
1757        assert!(auth.is_none());
1758        assert_eq!(client.subscription_args.len(), 1);
1759        let args = client
1760            .subscription_args
1761            .get("market_stats:0")
1762            .expect("replay arguments");
1763        assert_eq!(args.channel, channel);
1764        assert!(args.auth.is_none());
1765        assert_eq!(args.generation, 1);
1766    }
1767
1768    #[rstest]
1769    #[case::deltas(false)]
1770    #[case::depth(true)]
1771    #[tokio::test]
1772    async fn book_registration_rejects_previous_session(#[case] depth: bool) {
1773        let client = LighterWebSocketClient::new(
1774            Some("wss://example/test".to_string()),
1775            LighterEnvironment::Testnet,
1776            Arc::new(MarketRegistry::new()),
1777            TransportBackend::default(),
1778            30,
1779            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1780            None,
1781        );
1782        let market_index = 7;
1783        let generation = 42;
1784        let channel = LighterWsChannel::OrderBook(market_index);
1785        let topic = channel.topic_key();
1786
1787        let consumer = |subscribed| {
1788            if depth {
1789                HandlerCommand::SetDepthSub {
1790                    market_index,
1791                    subscribed,
1792                }
1793            } else {
1794                HandlerCommand::SetBookDeltasSub {
1795                    market_index,
1796                    subscribed,
1797                }
1798            }
1799        };
1800
1801        let (previous_tx, mut previous_rx) = tokio::sync::mpsc::unbounded_channel();
1802        *client.cmd_tx.write().await = previous_tx;
1803        let session = client
1804            .send_cmd_session(consumer(true))
1805            .await
1806            .expect("previous-session consumer registration");
1807        let (current_tx, mut current_rx) = tokio::sync::mpsc::unbounded_channel();
1808        *client.cmd_tx.write().await = current_tx;
1809        client.subscriptions.add_reference(&topic);
1810        client.subscriptions.mark_subscribe(&topic);
1811        client.subscriptions.confirm_subscribe(&topic);
1812        client
1813            .next_subscription_generation
1814            .store(generation + 1, Ordering::Relaxed);
1815        client.subscription_args.insert(
1816            topic.clone(),
1817            SubscriptionArgs {
1818                channel: channel.clone(),
1819                auth: None,
1820                generation,
1821            },
1822        );
1823
1824        let error = client
1825            .subscribe_order_book_stream(market_index, &session)
1826            .await
1827            .expect_err("previous-session book registration must fail");
1828        client
1829            .send_cmd_on_session(consumer(false), &session)
1830            .await
1831            .expect("skip previous-session rollback");
1832
1833        let LighterWsError::Client(message) = error else {
1834            panic!("expected session mismatch client error");
1835        };
1836
1837        assert_eq!(
1838            message,
1839            format!("WebSocket session changed before subscription for {topic}")
1840        );
1841        assert_eq!(client.subscriptions.get_reference_count(&topic), 1);
1842        assert_eq!(client.subscription_count(), 1);
1843        assert_eq!(client.subscription_args.len(), 1);
1844        assert_eq!(
1845            client.next_subscription_generation.load(Ordering::Relaxed),
1846            generation + 1
1847        );
1848        let args = client
1849            .subscription_args
1850            .get(&topic)
1851            .expect("current-session replay args");
1852        assert_eq!(args.channel, channel);
1853        assert!(args.auth.is_none());
1854        assert_eq!(args.generation, generation);
1855        let command = previous_rx
1856            .try_recv()
1857            .expect("previous-session consumer flag");
1858
1859        match (depth, command) {
1860            (
1861                true,
1862                HandlerCommand::SetDepthSub {
1863                    market_index: actual,
1864                    subscribed,
1865                },
1866            )
1867            | (
1868                false,
1869                HandlerCommand::SetBookDeltasSub {
1870                    market_index: actual,
1871                    subscribed,
1872                },
1873            ) => {
1874                assert_eq!(actual, market_index);
1875                assert!(subscribed);
1876            }
1877            (_, command) => panic!("unexpected previous-session command: {command:?}"),
1878        }
1879
1880        assert!(matches!(
1881            previous_rx.try_recv(),
1882            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1883        ));
1884        assert!(matches!(
1885            current_rx.try_recv(),
1886            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1887        ));
1888    }
1889
1890    #[tokio::test]
1891    async fn subscribe_waits_for_venue_result_and_removes_failed_generation() {
1892        let client = LighterWebSocketClient::new(
1893            Some("wss://example/test".to_string()),
1894            LighterEnvironment::Testnet,
1895            Arc::new(MarketRegistry::new()),
1896            TransportBackend::default(),
1897            30,
1898            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1899            None,
1900        );
1901        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1902        *client.cmd_tx.write().await = cmd_tx;
1903
1904        let subscribe_client = client.clone();
1905        let subscribe = get_runtime().spawn(async move {
1906            subscribe_client
1907                .subscribe_market_stats(LighterMarketSelection::Market(0))
1908                .await
1909        });
1910
1911        let command = cmd_rx.recv().await.expect("subscribe command");
1912        let HandlerCommand::Subscribe {
1913            response_tx: Some(response_tx),
1914            ..
1915        } = command
1916        else {
1917            panic!("expected subscribe command with venue result sender");
1918        };
1919
1920        assert!(!subscribe.is_finished());
1921        assert!(client.subscription_args.contains_key("market_stats:0"));
1922
1923        response_tx
1924            .send(Err("venue rejected subscription".to_string()))
1925            .expect("subscription result receiver");
1926        let error = subscribe
1927            .await
1928            .expect("subscribe task")
1929            .expect_err("failed venue open must fail the caller");
1930
1931        assert!(error.to_string().contains("venue rejected subscription"));
1932        assert!(!client.subscription_args.contains_key("market_stats:0"));
1933    }
1934
1935    #[tokio::test]
1936    async fn failed_older_subscribe_does_not_remove_newer_generation() {
1937        let client = LighterWebSocketClient::new(
1938            Some("wss://example/test".to_string()),
1939            LighterEnvironment::Testnet,
1940            Arc::new(MarketRegistry::new()),
1941            TransportBackend::default(),
1942            30,
1943            Duration::from_secs(DEFAULT_BOOK_SNAPSHOT_TIMEOUT_SECS),
1944            None,
1945        );
1946        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1947        *client.cmd_tx.write().await = cmd_tx;
1948
1949        let older_client = client.clone();
1950        let older = get_runtime().spawn(async move {
1951            older_client
1952                .subscribe_market_stats(LighterMarketSelection::Market(0))
1953                .await
1954        });
1955        let HandlerCommand::Subscribe {
1956            response_tx: Some(older_response),
1957            ..
1958        } = cmd_rx.recv().await.expect("older subscribe command")
1959        else {
1960            panic!("expected older subscribe command with venue result sender");
1961        };
1962
1963        let newer_client = client.clone();
1964        let newer = get_runtime().spawn(async move {
1965            newer_client
1966                .subscribe_market_stats(LighterMarketSelection::Market(0))
1967                .await
1968        });
1969        let HandlerCommand::Subscribe {
1970            response_tx: Some(newer_response),
1971            ..
1972        } = cmd_rx.recv().await.expect("newer subscribe command")
1973        else {
1974            panic!("expected newer subscribe command with venue result sender");
1975        };
1976
1977        older_response
1978            .send(Err("older generation failed".to_string()))
1979            .expect("older result receiver");
1980        older
1981            .await
1982            .expect("older subscribe task")
1983            .expect_err("older generation must fail");
1984        assert!(client.subscription_args.contains_key("market_stats:0"));
1985
1986        newer_response
1987            .send(Err("newer generation failed".to_string()))
1988            .expect("newer result receiver");
1989        newer
1990            .await
1991            .expect("newer subscribe task")
1992            .expect_err("newer generation must fail");
1993        assert!(!client.subscription_args.contains_key("market_stats:0"));
1994    }
1995
1996    fn stub_instrument(id: InstrumentId) -> InstrumentAny {
1997        InstrumentAny::CryptoPerpetual(
1998            CryptoPerpetual::builder()
1999                .instrument_id(id)
2000                .raw_symbol(id.symbol)
2001                .base_currency(Currency::from("ETH"))
2002                .quote_currency(Currency::from("USDC"))
2003                .settlement_currency(Currency::from("USDC"))
2004                .is_inverse(false)
2005                .price_precision(2)
2006                .size_precision(4)
2007                .price_increment(Price::from("0.01"))
2008                .size_increment(Quantity::from("0.0001"))
2009                .ts_event(UnixNanos::default())
2010                .ts_init(UnixNanos::default())
2011                .build()
2012                .unwrap(),
2013        )
2014    }
2015}
nautilus_lighter/websocket/client.rs

nautilus_lighter/websocket/
client.rs
