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

nautilus_lighter/websocket/
client.rs
