nautilus_bybit/websocket/
handler.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//! WebSocket message handler for Bybit.
17
18use std::{
19    collections::VecDeque,
20    sync::{
21        Arc,
22        atomic::{AtomicBool, AtomicU64, Ordering},
23    },
24};
25
26use dashmap::DashMap;
27use nautilus_core::string::secret::SecretString;
28use nautilus_live::book::snapshot::SnapshotGate;
29use nautilus_network::{
30    error::SendError,
31    retry::{RetryManager, create_websocket_retry_manager},
32    websocket::{AuthTracker, SubscriptionState, WebSocketClient},
33};
34use serde::Serialize;
35use serde_json::Value;
36use tokio_tungstenite::tungstenite::Message;
37use tokio_util::sync::CancellationToken;
38use ustr::Ustr;
39
40use super::{
41    enums::BybitWsOperation,
42    error::{BybitWsError, BybitWsResult, create_bybit_timeout_error, should_retry_bybit_error},
43    messages::{
44        BybitSubscription, BybitWebSocketError, BybitWsFrame, BybitWsMessage, BybitWsOrderResponse,
45        BybitWsResponse, BybitWsSubscriptionMsg,
46    },
47    parse::parse_bybit_ws_frame,
48};
49use crate::common::{
50    enums::{BybitProductType, BybitWsOrderRequestOp},
51    rate_limit::{
52        BYBIT_RATE_LIMIT_HEADER, BYBIT_RATE_LIMIT_RESET_HEADER, BYBIT_RATE_LIMIT_STATUS_HEADER,
53        BybitRateLimiter,
54    },
55};
56
57/// Semantic order command whose time-sensitive header is built at send time.
58#[derive(Debug)]
59pub struct BybitWsOrderCommand {
60    pub(crate) req_id: String,
61    pub(crate) op: BybitWsOrderRequestOp,
62    pub(crate) category: BybitProductType,
63    pub(crate) weight: u32,
64    pub(crate) referer: Option<String>,
65    pub(crate) args: Vec<Value>,
66}
67
68#[derive(Debug, Serialize)]
69#[serde(rename_all = "camelCase")]
70struct BybitWsSignedOrderRequest {
71    req_id: String,
72    op: BybitWsOrderRequestOp,
73    header: BybitWsSignedOrderHeader,
74    args: Vec<Value>,
75}
76
77#[derive(Debug, Serialize)]
78#[serde(rename_all = "SCREAMING-KEBAB-CASE")]
79struct BybitWsSignedOrderHeader {
80    x_bapi_timestamp: String,
81    x_bapi_recv_window: String,
82    #[serde(rename = "Referer", skip_serializing_if = "Option::is_none")]
83    referer: Option<String>,
84}
85
86#[derive(Clone, Copy, Debug)]
87struct PendingRate {
88    endpoint: &'static str,
89    category: BybitProductType,
90}
91
92#[derive(Debug)]
93enum OrderSendFailure {
94    NotSent(BybitWsError),
95    Ambiguous(BybitWsError),
96}
97
98/// Commands sent from the outer client to the inner message handler.
99#[derive(Debug)]
100pub enum HandlerCommand {
101    SetClient(WebSocketClient),
102    Disconnect,
103    Authenticate {
104        payload: SecretString,
105    },
106    Subscribe {
107        topics: Vec<String>,
108    },
109    /// Subscribes to a book topic, opening `gate` once the connection confirms the write.
110    ///
111    /// A queued write always runs, since another subscription can hold the same topic.
112    SubscribeBook {
113        payload: String,
114        gate: SnapshotGate,
115        completion: tokio::sync::oneshot::Sender<BybitWsResult<()>>,
116    },
117    Unsubscribe {
118        topics: Vec<String>,
119    },
120    /// Replaces a book subscription on the current connection, opening `gate` once both writes
121    /// complete.
122    ResubscribeBook {
123        topic: String,
124        cancel: CancellationToken,
125        gate: SnapshotGate,
126        completion: tokio::sync::oneshot::Sender<BybitWsResult<()>>,
127    },
128    SendOrder {
129        command: BybitWsOrderCommand,
130    },
131    SendOrders {
132        commands: Vec<BybitWsOrderCommand>,
133    },
134}
135
136pub(super) struct BybitWsFeedHandler {
137    signal: Arc<AtomicBool>,
138    inner: Option<WebSocketClient>,
139    cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
140    raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
141    auth_tracker: AuthTracker,
142    subscriptions: SubscriptionState,
143    rate_limiter: BybitRateLimiter,
144    recv_window_ms: Arc<AtomicU64>,
145    pending_rates: DashMap<String, PendingRate>,
146    pending_orders: VecDeque<BybitWsOrderCommand>,
147    retry_manager: RetryManager<BybitWsError>,
148}
149
150impl BybitWsFeedHandler {
151    /// Creates a new [`BybitWsFeedHandler`] instance.
152    pub(super) fn new(
153        signal: Arc<AtomicBool>,
154        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
155        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
156        auth_tracker: AuthTracker,
157        subscriptions: SubscriptionState,
158        rate_limiter: BybitRateLimiter,
159        recv_window_ms: Arc<AtomicU64>,
160    ) -> Self {
161        Self {
162            signal,
163            inner: None,
164            cmd_rx,
165            raw_rx,
166            auth_tracker,
167            subscriptions,
168            rate_limiter,
169            recv_window_ms,
170            pending_rates: DashMap::new(),
171            pending_orders: VecDeque::new(),
172            retry_manager: create_websocket_retry_manager(),
173        }
174    }
175
176    pub(super) fn is_stopped(&self) -> bool {
177        self.signal.load(Ordering::Relaxed)
178    }
179
180    /// Sends a WebSocket message with retry logic.
181    async fn send_with_retry(&self, payload: String) -> Result<(), BybitWsError> {
182        self.send_secret_with_retry(payload.into(), None).await
183    }
184
185    async fn send_secret_with_retry(
186        &self,
187        payload: SecretString,
188        connection_epoch: Option<u64>,
189    ) -> Result<(), BybitWsError> {
190        if let Some(client) = &self.inner {
191            self.retry_manager
192                .invocation(
193                    "websocket_send",
194                    || {
195                        let payload = payload.clone();
196                        async move {
197                            let payload = payload.expose_secret().to_owned();
198
199                            let result = match connection_epoch {
200                                Some(epoch) => {
201                                    client.send_text_on_connection(payload, None, epoch).await
202                                }
203                                None => client.send_text(payload, None).await,
204                            };
205
206                            result.map_err(|e| match e {
207                                SendError::ConnectionChanged => BybitWsError::Authentication(
208                                    "Connection changed before authentication".to_string(),
209                                ),
210                                e => BybitWsError::Transport(format!("Send failed: {e}")),
211                            })
212                        }
213                    },
214                    should_retry_bybit_error,
215                    |e| create_bybit_timeout_error(e.to_string()),
216                )
217                .execute()
218                .await
219        } else {
220            Err(BybitWsError::ClientError(
221                "No active WebSocket client".to_string(),
222            ))
223        }
224    }
225
226    // Confirms the write on the current connection; a failed write is not replayed on the next
227    async fn write_book_subscription(&self, payload: String) -> BybitWsResult<()> {
228        let client = self.inner.as_ref().ok_or(BybitWsError::NotConnected)?;
229        let epoch = client.connection_epoch();
230
231        client
232            .send_text_on_connection(payload, None, epoch)
233            .await
234            .map_err(|e| BybitWsError::Transport(e.to_string()))
235    }
236
237    // Writes both requests on one connection so the venue replaces the subscription it serves,
238    // and never stops between them, since a depth-1 topic can also carry quotes
239    async fn resubscribe_book(&self, topic: &str) -> BybitWsResult<()> {
240        let client = self.inner.as_ref().ok_or(BybitWsError::NotConnected)?;
241        let epoch = client.connection_epoch();
242
243        // Keeps the topic replayed on reconnect until the venue confirms the new subscription
244        self.subscriptions.mark_failure(topic);
245
246        let topic = topic.to_string();
247
248        for op in [BybitWsOperation::Unsubscribe, BybitWsOperation::Subscribe] {
249            let request = BybitSubscription {
250                op,
251                args: vec![topic.clone()],
252                req_id: Some(topic.clone()),
253            };
254
255            let payload = serde_json::to_string(&request)?;
256            client.send_text_on_connection(payload, None, epoch).await?;
257        }
258
259        Ok(())
260    }
261
262    async fn send_order(&self, command: &BybitWsOrderCommand) -> Result<bool, OrderSendFailure> {
263        if !self.auth_tracker.is_authenticated() {
264            return Ok(false);
265        }
266
267        let client = self.inner.as_ref().ok_or_else(|| {
268            OrderSendFailure::NotSent(BybitWsError::ClientError(
269                "No active WebSocket client".to_string(),
270            ))
271        })?;
272        let (endpoint, _) = order_operation(command.op);
273
274        self.rate_limiter
275            .acquire_ws_order(endpoint, command.category, command.weight)
276            .await
277            .map_err(|e| OrderSendFailure::NotSent(BybitWsError::ClientError(e)))?;
278
279        if !self.auth_tracker.is_authenticated() {
280            return Ok(false);
281        }
282
283        let connection_epoch = client.connection_epoch();
284        let recv_window_ms = self.recv_window_ms.load(Ordering::Acquire);
285        let request = BybitWsSignedOrderRequest {
286            req_id: command.req_id.clone(),
287            op: command.op,
288            header: BybitWsSignedOrderHeader {
289                x_bapi_timestamp: nautilus_core::time::get_atomic_clock_realtime()
290                    .get_time_ms()
291                    .to_string(),
292                x_bapi_recv_window: recv_window_ms.to_string(),
293                referer: command.referer.clone(),
294            },
295            args: command.args.clone(),
296        };
297        let payload = serde_json::to_string(&request)
298            .map_err(|e| OrderSendFailure::NotSent(BybitWsError::Json(e.to_string())))?;
299        self.pending_rates.insert(
300            command.req_id.clone(),
301            PendingRate {
302                endpoint,
303                category: command.category,
304            },
305        );
306
307        match client
308            .send_text_on_connection(payload, None, connection_epoch)
309            .await
310        {
311            Ok(()) => Ok(true),
312            Err(e) => self.classify_order_send_error(&command.req_id, e),
313        }
314    }
315
316    fn classify_order_send_error(
317        &self,
318        req_id: &str,
319        error: SendError,
320    ) -> Result<bool, OrderSendFailure> {
321        match error {
322            SendError::ConnectionChanged => {
323                self.pending_rates.remove(req_id);
324                self.auth_tracker.invalidate();
325                Ok(false)
326            }
327            SendError::Timeout => {
328                self.pending_rates.remove(req_id);
329                self.auth_tracker.invalidate();
330                Err(OrderSendFailure::NotSent(BybitWsError::ClientError(
331                    "Order command was not written".to_string(),
332                )))
333            }
334            SendError::InvalidInput(_) | SendError::BufferFull | SendError::Closed => {
335                self.pending_rates.remove(req_id);
336                Err(OrderSendFailure::NotSent(BybitWsError::ClientError(
337                    "Order command was not written".to_string(),
338                )))
339            }
340            error => Err(OrderSendFailure::Ambiguous(BybitWsError::Send(
341                error.to_string(),
342            ))),
343        }
344    }
345
346    pub(super) async fn next(&mut self) -> Option<BybitWsMessage> {
347        loop {
348            if self.auth_tracker.is_authenticated()
349                && let Some(command) = self.pending_orders.pop_front()
350            {
351                let req_id = command.req_id.clone();
352                match self.send_order(&command).await {
353                    Ok(true) => {}
354                    Ok(false) => self.pending_orders.push_front(command),
355                    Err(OrderSendFailure::NotSent(error)) => {
356                        return Some(order_not_sent_message(&command, &error));
357                    }
358                    Err(OrderSendFailure::Ambiguous(error)) => {
359                        log::warn!("Ambiguous order send failure: req_id={req_id}, error={error}");
360                    }
361                }
362                continue;
363            }
364
365            tokio::select! {
366                Some(cmd) = self.cmd_rx.recv() => {
367                    match cmd {
368                        HandlerCommand::SetClient(client) => {
369                            log::debug!("WebSocketClient received by handler");
370                            self.inner = Some(client);
371                        }
372                        HandlerCommand::Disconnect => {
373                            log::debug!("Disconnect command received");
374
375                            if let Some(client) = self.inner.take() {
376                                client.disconnect().await;
377                            }
378                        }
379                        HandlerCommand::Authenticate { payload } => {
380                            log::debug!("Authenticate command received");
381
382                            let epoch = self.inner.as_ref().map(WebSocketClient::connection_epoch);
383                            if let Err(e) = self.send_secret_with_retry(payload, epoch).await {
384                                log::error!("Failed to send authentication after retries: {e}");
385                            }
386                        }
387                        HandlerCommand::Subscribe { topics } => {
388                            for topic in topics {
389                                log::debug!("Subscribing to topic: topic={topic}");
390                                if let Err(e) = self.send_with_retry(topic.clone()).await {
391                                    log::error!("Failed to send subscription after retries: topic={topic}, error={e}");
392                                }
393                            }
394                        }
395                        HandlerCommand::SubscribeBook { payload, gate, completion } => {
396                            let result = self.write_book_subscription(payload).await;
397
398                            if result.is_ok() {
399                                gate.open();
400                            }
401                            let _ = completion.send(result);
402                        }
403                        HandlerCommand::Unsubscribe { topics } => {
404                            for topic in topics {
405                                log::debug!("Unsubscribing from topic: topic={topic}");
406                                if let Err(e) = self.send_with_retry(topic.clone()).await {
407                                    log::error!("Failed to send unsubscription after retries: topic={topic}, error={e}");
408                                }
409                            }
410                        }
411                        HandlerCommand::ResubscribeBook { topic, cancel, gate, completion } => {
412                            if cancel.is_cancelled() {
413                                continue;
414                            }
415
416                            let result = self.resubscribe_book(&topic).await;
417
418                            if result.is_ok() {
419                                gate.open();
420                            }
421                            let _ = completion.send(result);
422                        }
423                        HandlerCommand::SendOrder { command } => {
424                            let req_id = command.req_id.clone();
425                            match self.send_order(&command).await {
426                                Ok(true) => {}
427                                Ok(false) => self.pending_orders.push_back(command),
428                                Err(OrderSendFailure::NotSent(error)) => {
429                                    return Some(order_not_sent_message(&command, &error));
430                                }
431                                Err(OrderSendFailure::Ambiguous(error)) => {
432                                    log::warn!(
433                                        "Ambiguous order send failure: req_id={req_id}, error={error}"
434                                    );
435                                }
436                            }
437                        }
438                        HandlerCommand::SendOrders { commands } => {
439                            self.pending_orders.extend(commands);
440                        }
441                    }
442                }
443
444                () = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
445                    if self.signal.load(Ordering::Relaxed) {
446                        log::debug!("Stop signal received during idle period");
447                        return None;
448                    }
449                }
450
451                msg = self.raw_rx.recv() => {
452                    let msg = match msg {
453                        Some(msg) => msg,
454                        None => {
455                            log::debug!("WebSocket stream closed");
456                            return None;
457                        }
458                    };
459
460                    if let Message::Ping(data) = &msg {
461                        log::trace!("Received ping frame with {} bytes", data.len());
462
463                        if let Some(client) = &self.inner
464                            && let Err(e) = client.send_pong(data.to_vec()).await
465                        {
466                            log::warn!("Failed to send pong frame: error={e}");
467                        }
468                        continue;
469                    }
470
471                    let frame = match Self::parse_raw_frame(msg) {
472                        Some(frame) => frame,
473                        None => continue,
474                    };
475
476                    if self.signal.load(Ordering::Relaxed) {
477                        log::debug!("Stop signal received");
478                        return None;
479                    }
480
481                    match frame {
482                        BybitWsFrame::Subscription(ref sub_msg) => {
483                            self.handle_subscription_ack(sub_msg);
484                        }
485                        BybitWsFrame::Auth(auth_response) => {
486                            let is_success = auth_response.success.unwrap_or(false)
487                                || (auth_response.ret_code == Some(0));
488
489                            if is_success {
490                                self.auth_tracker.succeed();
491                                log::debug!("WebSocket authenticated");
492                            } else {
493                                let error_msg = auth_response
494                                    .ret_msg
495                                    .as_deref()
496                                    .unwrap_or("Authentication rejected");
497                                self.auth_tracker.fail(error_msg);
498                                log::error!("WebSocket authentication failed: error={error_msg}");
499                            }
500                            return Some(BybitWsMessage::Auth(auth_response));
501                        }
502                        BybitWsFrame::ErrorResponse(ref resp) => {
503                            // Failed subscription/unsubscription ACKs arrive as
504                            // ErrorResponse when success=false. Route them through
505                            // the subscription state machine when the op field is
506                            // present, otherwise forward as a generic error.
507                            if let Some(op) = &resp.op {
508                                if *op == BybitWsOperation::Subscribe
509                                    || *op == BybitWsOperation::Unsubscribe
510                                {
511                                    self.handle_subscription_error(resp);
512                                } else {
513                                    let error = BybitWebSocketError::from_response(resp);
514                                    return Some(BybitWsMessage::Error(error));
515                                }
516                            } else {
517                                let error = BybitWebSocketError::from_response(resp);
518                                return Some(BybitWsMessage::Error(error));
519                            }
520                        }
521                        BybitWsFrame::OrderResponse(resp) => {
522                            self.observe_order_rate_limit(&resp);
523                            return Some(BybitWsMessage::OrderResponse(resp));
524                        }
525                        BybitWsFrame::Orderbook(msg) => {
526                            return Some(BybitWsMessage::Orderbook(msg));
527                        }
528                        BybitWsFrame::Trade(msg) => {
529                            return Some(BybitWsMessage::Trade(msg));
530                        }
531                        BybitWsFrame::Liquidation(msg) => {
532                            return Some(BybitWsMessage::Liquidation(msg));
533                        }
534                        BybitWsFrame::Kline(msg) => {
535                            return Some(BybitWsMessage::Kline(msg));
536                        }
537                        BybitWsFrame::TickerLinear(msg) => {
538                            return Some(BybitWsMessage::TickerLinear(msg));
539                        }
540                        BybitWsFrame::TickerOption(msg) => {
541                            return Some(BybitWsMessage::TickerOption(msg));
542                        }
543                        BybitWsFrame::AccountOrder(msg) => {
544                            return Some(BybitWsMessage::AccountOrder(msg));
545                        }
546                        BybitWsFrame::AccountExecution(msg) => {
547                            return Some(BybitWsMessage::AccountExecution(msg));
548                        }
549                        BybitWsFrame::AccountExecutionFast(msg) => {
550                            return Some(BybitWsMessage::AccountExecutionFast(msg));
551                        }
552                        BybitWsFrame::AccountWallet(msg) => {
553                            return Some(BybitWsMessage::AccountWallet(msg));
554                        }
555                        BybitWsFrame::AccountPosition(msg) => {
556                            return Some(BybitWsMessage::AccountPosition(msg));
557                        }
558                        BybitWsFrame::Reconnected => {
559                            self.auth_tracker.invalidate();
560                            return Some(BybitWsMessage::Reconnected);
561                        }
562                        BybitWsFrame::Unknown(value) => {
563                            log::debug!("Unknown WebSocket frame: {value}");
564                        }
565                    }
566                }
567            }
568        }
569    }
570
571    fn observe_order_rate_limit(&self, response: &super::messages::BybitWsOrderResponse) {
572        let Some(req_id) = response.req_id.as_deref() else {
573            return;
574        };
575        let Some((_, pending)) = self.pending_rates.remove(req_id) else {
576            return;
577        };
578        let Some(header) = response.header.as_ref() else {
579            return;
580        };
581        let parse_u32 = |key: &str| {
582            header.get(key).and_then(|value| {
583                value
584                    .as_u64()
585                    .and_then(|value| u32::try_from(value).ok())
586                    .or_else(|| value.as_str()?.parse::<u32>().ok())
587            })
588        };
589        let Some(limit) = parse_u32(BYBIT_RATE_LIMIT_HEADER) else {
590            return;
591        };
592        let Some(remaining) = parse_u32(BYBIT_RATE_LIMIT_STATUS_HEADER) else {
593            return;
594        };
595        let reset_timestamp_ms = header.get(BYBIT_RATE_LIMIT_RESET_HEADER).and_then(|value| {
596            value
597                .as_i64()
598                .or_else(|| value.as_str()?.parse::<i64>().ok())
599        });
600        self.rate_limiter.observe_account(
601            pending.endpoint,
602            Some(pending.category),
603            limit,
604            remaining,
605            reset_timestamp_ms,
606        );
607    }
608
609    fn handle_subscription_ack(&self, sub_msg: &BybitWsSubscriptionMsg) {
610        match sub_msg.op {
611            BybitWsOperation::Subscribe => {
612                if sub_msg.success {
613                    if let Some(topic) = &sub_msg.req_id {
614                        self.subscriptions.confirm_subscribe(topic);
615                        log::debug!("Subscription confirmed: topic={topic}");
616                    } else {
617                        // No req_id, fall back to confirming all pending
618                        for topic in self.subscriptions.pending_subscribe_topics() {
619                            self.subscriptions.confirm_subscribe(&topic);
620                            log::debug!("Subscription confirmed (bulk): topic={topic}");
621                        }
622                    }
623                } else if let Some(topic) = &sub_msg.req_id {
624                    self.subscriptions.mark_failure(topic);
625                    log::warn!(
626                        "Subscription failed: topic={topic}, error={:?}",
627                        sub_msg.ret_msg
628                    );
629                } else {
630                    for topic in self.subscriptions.pending_subscribe_topics() {
631                        self.subscriptions.mark_failure(&topic);
632                        log::warn!(
633                            "Subscription failed (bulk): topic={topic}, error={:?}",
634                            sub_msg.ret_msg
635                        );
636                    }
637                }
638            }
639            BybitWsOperation::Unsubscribe => {
640                if sub_msg.success {
641                    if let Some(topic) = &sub_msg.req_id {
642                        self.subscriptions.confirm_unsubscribe(topic);
643                        log::debug!("Unsubscription confirmed: topic={topic}");
644                    } else {
645                        for topic in self.subscriptions.pending_unsubscribe_topics() {
646                            self.subscriptions.confirm_unsubscribe(&topic);
647                            log::debug!("Unsubscription confirmed (bulk): topic={topic}");
648                        }
649                    }
650                } else {
651                    let topic_desc = sub_msg.req_id.as_deref().unwrap_or("unknown");
652                    log::warn!(
653                        "Unsubscription failed: topic={topic_desc}, error={:?}",
654                        sub_msg.ret_msg
655                    );
656                }
657            }
658            _ => {}
659        }
660    }
661
662    fn handle_subscription_error(&self, resp: &BybitWsResponse) {
663        let topic = resp.req_id.as_deref().unwrap_or("unknown");
664        let error_msg = resp.ret_msg.as_deref().unwrap_or("unknown error");
665
666        match resp.op {
667            Some(BybitWsOperation::Subscribe) => {
668                // Duplicate subscribe is harmless: the topic is active on the
669                // venue, so confirm it instead of looping retries every reconnect.
670                if is_already_subscribed_error(error_msg)
671                    && let Some(ref req_id) = resp.req_id
672                {
673                    self.subscriptions.confirm_subscribe(req_id);
674                    log::debug!("Subscription duplicate ignored: topic={topic}, error={error_msg}");
675                    return;
676                }
677
678                if let Some(ref req_id) = resp.req_id {
679                    self.subscriptions.mark_failure(req_id);
680                } else {
681                    for t in self.subscriptions.pending_subscribe_topics() {
682                        self.subscriptions.mark_failure(&t);
683                    }
684                }
685                log::warn!("Subscription error: topic={topic}, error={error_msg}");
686            }
687            Some(BybitWsOperation::Unsubscribe) => {
688                log::warn!("Unsubscription error: topic={topic}, error={error_msg}");
689            }
690            _ => {}
691        }
692    }
693
694    fn parse_raw_frame(msg: Message) -> Option<BybitWsFrame> {
695        match msg {
696            Message::Text(text) => {
697                if text == nautilus_network::RECONNECTED {
698                    log::debug!("Received WebSocket reconnected signal");
699                    return Some(BybitWsFrame::Reconnected);
700                }
701
702                if text.trim().eq_ignore_ascii_case("pong") {
703                    return None;
704                }
705
706                log::trace!("Raw websocket message: {text}");
707
708                let value: serde_json::Value = match serde_json::from_str(&text) {
709                    Ok(v) => v,
710                    Err(e) => {
711                        log::error!("Failed to parse WebSocket message: {e}: {text}");
712                        return None;
713                    }
714                };
715
716                if value
717                    .get("op")
718                    .and_then(serde_json::Value::as_str)
719                    .is_some_and(|op| op == BybitWsOperation::Pong.as_ref())
720                {
721                    return None;
722                }
723
724                Some(parse_bybit_ws_frame(value))
725            }
726            Message::Binary(msg) => {
727                log::debug!("Raw binary frame ({} bytes)", msg.len());
728                log::trace!("Raw binary: {msg:?}");
729                None
730            }
731            Message::Close(_) => {
732                log::debug!("Received close message, waiting for reconnection");
733                None
734            }
735            _ => None,
736        }
737    }
738}
739
740fn order_operation(op: BybitWsOrderRequestOp) -> (&'static str, &'static str) {
741    match op {
742        BybitWsOrderRequestOp::Create => ("/v5/order/create", "order.create"),
743        BybitWsOrderRequestOp::Amend => ("/v5/order/amend", "order.amend"),
744        BybitWsOrderRequestOp::Cancel => ("/v5/order/cancel", "order.cancel"),
745        BybitWsOrderRequestOp::CreateBatch => ("/v5/order/create-batch", "order.create-batch"),
746        BybitWsOrderRequestOp::AmendBatch => ("/v5/order/amend-batch", "order.amend-batch"),
747        BybitWsOrderRequestOp::CancelBatch => ("/v5/order/cancel-batch", "order.cancel-batch"),
748    }
749}
750
751fn order_not_sent_message(command: &BybitWsOrderCommand, error: &BybitWsError) -> BybitWsMessage {
752    let (_, op) = order_operation(command.op);
753    BybitWsMessage::OrderResponse(BybitWsOrderResponse {
754        op: Ustr::from(op),
755        conn_id: None,
756        ret_code: -1,
757        ret_msg: error.to_string(),
758        data: Value::Object(serde_json::Map::new()),
759        req_id: Some(command.req_id.clone()),
760        header: None,
761        ret_ext_info: None,
762    })
763}
764
765fn is_already_subscribed_error(error_msg: &str) -> bool {
766    error_msg
767        .to_ascii_lowercase()
768        .contains("already subscribed")
769}
770
771#[cfg(test)]
772mod tests {
773    use std::time::Duration;
774
775    use futures_util::{SinkExt, StreamExt};
776    use nautilus_common::testing::wait_until_async;
777    use nautilus_core::string::secret::REDACTED;
778    use nautilus_network::websocket::{WebSocketConfig, channel_message_handler};
779    use rstest::rstest;
780    use ustr::Ustr;
781
782    use super::*;
783    use crate::common::{
784        consts::BYBIT_WS_TOPIC_DELIMITER, rate_limit::BybitRateLimiter, testing::load_test_json,
785    };
786
787    fn create_test_handler() -> BybitWsFeedHandler {
788        let signal = Arc::new(AtomicBool::new(false));
789        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
790        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
791        let auth_tracker = AuthTracker::new();
792        let subscriptions = SubscriptionState::new(BYBIT_WS_TOPIC_DELIMITER);
793
794        BybitWsFeedHandler::new(
795            signal,
796            cmd_rx,
797            raw_rx,
798            auth_tracker,
799            subscriptions,
800            BybitRateLimiter::for_websocket(
801                "wss://bybit-handler-test.invalid",
802                Some("test-key"),
803                None,
804            ),
805            Arc::new(AtomicU64::new(5_000)),
806        )
807    }
808
809    #[tokio::test]
810    async fn test_authentication_with_full_replay_buffer() {
811        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
812        let address = listener.local_addr().unwrap();
813        let (replayed_tx, replayed_rx) = tokio::sync::oneshot::channel();
814
815        let server = tokio::spawn(async move {
816            let (first, _) = listener.accept().await.unwrap();
817            let _first = tokio_tungstenite::accept_async(first).await.unwrap();
818            let (replacement, _) = listener.accept().await.unwrap();
819            let mut replacement = tokio_tungstenite::accept_async(replacement).await.unwrap();
820            assert_eq!(
821                replacement.next().await.unwrap().unwrap(),
822                Message::text("authenticate")
823            );
824            replacement
825                .send(Message::text(load_test_json("ws_auth_success.json")))
826                .await
827                .unwrap();
828            assert_eq!(
829                replacement.next().await.unwrap().unwrap(),
830                Message::text("held")
831            );
832            replayed_tx.send(()).unwrap();
833            std::future::pending::<()>().await;
834        });
835
836        let mut handler = create_test_handler();
837        let tracker = handler.auth_tracker.clone();
838        let (message_handler, raw_rx) = channel_message_handler();
839        handler.raw_rx = raw_rx;
840        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
841        handler.cmd_rx = cmd_rx;
842        let config = WebSocketConfig::builder()
843            .url(format!("ws://{address}"))
844            .writer_capacity(1)
845            .reconnect_delay_initial_ms(1)
846            .reconnect_delay_max_ms(1)
847            .reconnect_jitter_ms(0)
848            .build()
849            .unwrap();
850        let client = WebSocketClient::builder()
851            .config(config)
852            .message_handler(message_handler)
853            .connect()
854            .await
855            .unwrap();
856        client.set_auth_tracker(tracker.clone(), true);
857
858        // Enqueue and request reconnect without yielding so the writer retains this message
859        client.send_text("held".to_string(), None).await.unwrap();
860        assert!(client.request_reconnect());
861        wait_until_async(
862            || async { client.is_active() && client.connection_epoch() == 1 },
863            Duration::from_secs(5),
864        )
865        .await;
866        assert!(matches!(
867            client.send_text("overflow".to_string(), None).await,
868            Err(SendError::BufferFull)
869        ));
870        handler.inner = Some(client);
871        let stale = tokio::time::timeout(
872            Duration::from_secs(5),
873            handler.send_secret_with_retry(SecretString::from("stale-auth".to_string()), Some(0)),
874        )
875        .await
876        .unwrap();
877        assert!(
878            matches!(stale, Err(BybitWsError::Authentication(message)) if message == "Connection changed before authentication")
879        );
880        cmd_tx
881            .send(HandlerCommand::Authenticate {
882                payload: SecretString::from("authenticate".to_string()),
883            })
884            .unwrap();
885
886        tokio::time::timeout(Duration::from_secs(5), async {
887            loop {
888                if let Some(BybitWsMessage::Auth(response)) = handler.next().await {
889                    assert_eq!(response.success, Some(true));
890                    break;
891                }
892            }
893        })
894        .await
895        .unwrap();
896
897        assert!(tracker.is_authenticated());
898        tokio::time::timeout(Duration::from_secs(5), replayed_rx)
899            .await
900            .unwrap()
901            .unwrap();
902        handler.inner.as_ref().unwrap().disconnect().await;
903        server.abort();
904    }
905
906    async fn connect_book_handler() -> (
907        BybitWsFeedHandler,
908        tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
909        tokio::sync::mpsc::UnboundedReceiver<Value>,
910    ) {
911        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
912        let address = listener.local_addr().unwrap();
913        let (requests_tx, requests_rx) = tokio::sync::mpsc::unbounded_channel();
914
915        // Forwards each request as it arrives, since a completed write may still be queued
916        tokio::spawn(async move {
917            let (stream, _) = listener.accept().await.unwrap();
918            let mut socket = tokio_tungstenite::accept_async(stream).await.unwrap();
919
920            while let Some(Ok(Message::Text(text))) = socket.next().await {
921                let request = serde_json::from_str::<Value>(&text).unwrap();
922
923                if requests_tx.send(request).is_err() {
924                    break;
925                }
926            }
927        });
928
929        let mut handler = create_test_handler();
930        let (message_handler, raw_rx) = channel_message_handler();
931        handler.raw_rx = raw_rx;
932        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
933        handler.cmd_rx = cmd_rx;
934        let config = WebSocketConfig::builder()
935            .url(format!("ws://{address}"))
936            .build()
937            .unwrap();
938        let client = WebSocketClient::builder()
939            .config(config)
940            .message_handler(message_handler)
941            .connect()
942            .await
943            .unwrap();
944        handler.inner = Some(client);
945
946        (handler, cmd_tx, requests_rx)
947    }
948
949    async fn complete_book_command(
950        handler: &mut BybitWsFeedHandler,
951        receiver: tokio::sync::oneshot::Receiver<BybitWsResult<()>>,
952    ) -> Result<BybitWsResult<()>, tokio::sync::oneshot::error::RecvError> {
953        let drive = async {
954            loop {
955                handler.next().await;
956            }
957        };
958
959        tokio::time::timeout(Duration::from_secs(5), async {
960            tokio::select! {
961                result = receiver => result,
962                () = drive => unreachable!("handler loop never ends"),
963            }
964        })
965        .await
966        .unwrap()
967    }
968
969    async fn next_request(requests: &mut tokio::sync::mpsc::UnboundedReceiver<Value>) -> Value {
970        tokio::time::timeout(Duration::from_secs(5), requests.recv())
971            .await
972            .unwrap()
973            .unwrap()
974    }
975
976    fn subscription_request(op: &str, topic: &str) -> Value {
977        serde_json::json!({"op": op, "args": [topic], "req_id": topic})
978    }
979
980    fn subscribe_book_command(
981        topic: &str,
982        gate: SnapshotGate,
983    ) -> (
984        HandlerCommand,
985        tokio::sync::oneshot::Receiver<BybitWsResult<()>>,
986    ) {
987        let payload = subscription_request("subscribe", topic).to_string();
988        let (completion, receiver) = tokio::sync::oneshot::channel();
989
990        let command = HandlerCommand::SubscribeBook {
991            payload,
992            gate,
993            completion,
994        };
995
996        (command, receiver)
997    }
998
999    #[tokio::test]
1000    async fn test_subscribe_book_opens_gate_after_write() {
1001        let (mut handler, cmd_tx, mut requests) = connect_book_handler().await;
1002        let topic = "orderbook.50.BTCUSDT";
1003        let gate = SnapshotGate::default();
1004        gate.lock().close();
1005        let (command, receiver) = subscribe_book_command(topic, gate.clone());
1006
1007        cmd_tx.send(command).unwrap();
1008        let result = complete_book_command(&mut handler, receiver).await;
1009        let request = next_request(&mut requests).await;
1010
1011        assert!(matches!(result, Ok(Ok(()))));
1012        assert!(!gate.lock().is_closed());
1013        assert_eq!(request, subscription_request("subscribe", topic));
1014        handler.inner.as_ref().unwrap().disconnect().await;
1015    }
1016
1017    #[tokio::test]
1018    async fn test_resubscribe_book_replaces_subscription_and_opens_gate() {
1019        let (mut handler, cmd_tx, mut requests) = connect_book_handler().await;
1020        let topic = "orderbook.50.BTCUSDT";
1021        handler.subscriptions.mark_subscribe(topic);
1022        handler.subscriptions.confirm_subscribe(topic);
1023        let gate = SnapshotGate::default();
1024        gate.lock().close();
1025        let (completion, receiver) = tokio::sync::oneshot::channel();
1026
1027        cmd_tx
1028            .send(HandlerCommand::ResubscribeBook {
1029                topic: topic.to_string(),
1030                cancel: CancellationToken::new(),
1031                gate: gate.clone(),
1032                completion,
1033            })
1034            .unwrap();
1035
1036        let result = complete_book_command(&mut handler, receiver).await;
1037        let pending = handler.subscriptions.pending_subscribe_topics();
1038        let unsubscribe = next_request(&mut requests).await;
1039        let subscribe = next_request(&mut requests).await;
1040
1041        assert!(matches!(result, Ok(Ok(()))));
1042        assert!(!gate.lock().is_closed());
1043        assert_eq!(pending, [topic]);
1044        assert_eq!(unsubscribe, subscription_request("unsubscribe", topic));
1045        assert_eq!(subscribe, subscription_request("subscribe", topic));
1046        handler.inner.as_ref().unwrap().disconnect().await;
1047    }
1048
1049    #[tokio::test]
1050    async fn test_subscribe_book_write_failure_keeps_gate_closed() {
1051        let (mut handler, cmd_tx, _requests) = connect_book_handler().await;
1052        handler.inner.as_ref().unwrap().disconnect().await;
1053        let gate = SnapshotGate::default();
1054        gate.lock().close();
1055        let (command, receiver) = subscribe_book_command("orderbook.50.BTCUSDT", gate.clone());
1056
1057        cmd_tx.send(command).unwrap();
1058        let result = complete_book_command(&mut handler, receiver).await;
1059
1060        assert!(matches!(result, Ok(Err(BybitWsError::Transport(_)))));
1061        assert!(gate.lock().is_closed());
1062    }
1063
1064    // A later subscribe arrives first; any replacement write would complete before its completion
1065    #[tokio::test]
1066    async fn test_resubscribe_book_cancelled_before_write_sends_nothing() {
1067        let (mut handler, cmd_tx, mut requests) = connect_book_handler().await;
1068        let topic = "orderbook.50.BTCUSDT";
1069        handler.subscriptions.mark_subscribe(topic);
1070        handler.subscriptions.confirm_subscribe(topic);
1071        let gate = SnapshotGate::default();
1072        gate.lock().close();
1073        let cancel = CancellationToken::new();
1074        cancel.cancel();
1075        let (completion, receiver) = tokio::sync::oneshot::channel();
1076        let sentinel = "orderbook.50.ETHUSDT";
1077        let (sentinel_command, sentinel_receiver) =
1078            subscribe_book_command(sentinel, SnapshotGate::default());
1079
1080        cmd_tx
1081            .send(HandlerCommand::ResubscribeBook {
1082                topic: topic.to_string(),
1083                cancel,
1084                gate: gate.clone(),
1085                completion,
1086            })
1087            .unwrap();
1088
1089        let result = complete_book_command(&mut handler, receiver).await;
1090        cmd_tx.send(sentinel_command).unwrap();
1091        let sentinel_result = complete_book_command(&mut handler, sentinel_receiver).await;
1092        let pending = handler.subscriptions.pending_subscribe_topics();
1093        let request = next_request(&mut requests).await;
1094
1095        assert!(result.is_err());
1096        assert!(matches!(sentinel_result, Ok(Ok(()))));
1097        assert!(gate.lock().is_closed());
1098        assert!(pending.is_empty());
1099        assert_eq!(request, subscription_request("subscribe", sentinel));
1100        handler.inner.as_ref().unwrap().disconnect().await;
1101    }
1102
1103    #[rstest]
1104    fn test_authenticate_command_debug_redacts_payload() {
1105        let payload = "authentication-secret";
1106        let command = HandlerCommand::Authenticate {
1107            payload: SecretString::from(payload.to_string()),
1108        };
1109
1110        let debug = format!("{command:?}");
1111
1112        assert!(debug.contains(REDACTED));
1113        assert!(!debug.contains(payload));
1114    }
1115
1116    fn load_value(fixture: &str) -> serde_json::Value {
1117        let json = load_test_json(fixture);
1118        serde_json::from_str(&json).unwrap()
1119    }
1120
1121    #[rstest]
1122    fn test_handler_initializes() {
1123        let _handler = create_test_handler();
1124    }
1125
1126    #[tokio::test]
1127    async fn order_without_writer_produces_correlated_terminal_response() {
1128        let handler = create_test_handler();
1129        handler.auth_tracker.succeed();
1130        let command = BybitWsOrderCommand {
1131            req_id: "not-sent-request".to_string(),
1132            op: BybitWsOrderRequestOp::Create,
1133            category: BybitProductType::Linear,
1134            weight: 1,
1135            referer: None,
1136            args: vec![serde_json::json!({"category": "linear"})],
1137        };
1138
1139        let failure = handler.send_order(&command).await.unwrap_err();
1140        let OrderSendFailure::NotSent(error) = failure else {
1141            panic!("Expected definitive not-sent failure, was {failure:?}");
1142        };
1143        let BybitWsMessage::OrderResponse(response) = order_not_sent_message(&command, &error)
1144        else {
1145            panic!("Expected order response");
1146        };
1147
1148        assert_eq!(response.op, "order.create");
1149        assert_eq!(response.ret_code, -1);
1150        assert_eq!(response.req_id.as_deref(), Some("not-sent-request"));
1151        assert_eq!(response.ret_msg, "Client error: No active WebSocket client");
1152    }
1153
1154    #[rstest]
1155    fn connection_change_invalidates_authentication_before_requeue() {
1156        let handler = create_test_handler();
1157        handler.auth_tracker.succeed();
1158        handler.pending_rates.insert(
1159            "reconnect-request".to_string(),
1160            PendingRate {
1161                endpoint: "/v5/order/create",
1162                category: BybitProductType::Linear,
1163            },
1164        );
1165
1166        let result = handler
1167            .classify_order_send_error("reconnect-request", SendError::ConnectionChanged)
1168            .unwrap();
1169
1170        assert!(!result);
1171        assert!(!handler.auth_tracker.is_authenticated());
1172        assert!(!handler.pending_rates.contains_key("reconnect-request"));
1173    }
1174
1175    #[rstest]
1176    fn pre_active_timeout_is_terminal_not_sent() {
1177        let handler = create_test_handler();
1178        handler.auth_tracker.succeed();
1179        handler.pending_rates.insert(
1180            "timeout-request".to_string(),
1181            PendingRate {
1182                endpoint: "/v5/order/create",
1183                category: BybitProductType::Linear,
1184            },
1185        );
1186
1187        let failure = handler
1188            .classify_order_send_error("timeout-request", SendError::Timeout)
1189            .unwrap_err();
1190
1191        assert!(matches!(failure, OrderSendFailure::NotSent(_)));
1192        assert!(!handler.auth_tracker.is_authenticated());
1193        assert!(!handler.pending_rates.contains_key("timeout-request"));
1194    }
1195
1196    #[rstest]
1197    fn buffer_full_is_not_sent_and_preserves_authentication() {
1198        let handler = create_test_handler();
1199        handler.auth_tracker.succeed();
1200        handler.pending_rates.insert(
1201            "full-request".to_string(),
1202            PendingRate {
1203                endpoint: "/v5/order/create",
1204                category: BybitProductType::Linear,
1205            },
1206        );
1207
1208        let failure = handler
1209            .classify_order_send_error("full-request", SendError::BufferFull)
1210            .unwrap_err();
1211
1212        assert!(matches!(failure, OrderSendFailure::NotSent(_)));
1213        assert!(handler.auth_tracker.is_authenticated());
1214        assert!(!handler.pending_rates.contains_key("full-request"));
1215    }
1216
1217    #[rstest]
1218    fn test_parse_frame_auth_success() {
1219        let value = load_value("ws_auth_success.json");
1220        let frame = parse_bybit_ws_frame(value);
1221        match frame {
1222            BybitWsFrame::Auth(auth) => {
1223                assert_eq!(auth.conn_id.as_deref(), Some("cejreaspqfm9se7usbrg-2xh"));
1224                assert_eq!(auth.ret_code, Some(0));
1225                assert_eq!(auth.success, Some(true));
1226            }
1227            other => panic!("Expected Auth, was {other:?}"),
1228        }
1229    }
1230
1231    #[rstest]
1232    fn test_parse_frame_auth_failure() {
1233        let value = load_value("ws_auth_failure.json");
1234        let frame = parse_bybit_ws_frame(value);
1235        match frame {
1236            BybitWsFrame::ErrorResponse(resp) => {
1237                assert_eq!(resp.ret_code, Some(10003));
1238                assert_eq!(resp.ret_msg.as_deref(), Some("Invalid apikey"));
1239            }
1240            other => panic!("Expected ErrorResponse, was {other:?}"),
1241        }
1242    }
1243
1244    #[rstest]
1245    fn test_parse_frame_subscription_ack() {
1246        let value = load_value("ws_subscription_ack.json");
1247        let frame = parse_bybit_ws_frame(value);
1248        match frame {
1249            BybitWsFrame::Subscription(sub) => {
1250                assert!(sub.success);
1251                assert_eq!(sub.op, BybitWsOperation::Subscribe);
1252                assert_eq!(sub.req_id.as_deref(), Some("sub-orderbook-1"));
1253            }
1254            other => panic!("Expected Subscription, was {other:?}"),
1255        }
1256    }
1257
1258    #[rstest]
1259    fn test_parse_frame_subscription_failure() {
1260        let value = load_value("ws_subscription_failure.json");
1261        let frame = parse_bybit_ws_frame(value);
1262        match frame {
1263            BybitWsFrame::ErrorResponse(resp) => {
1264                assert_eq!(
1265                    resp.ret_msg.as_deref(),
1266                    Some("Invalid topic: invalid.topic.BTCUSDT")
1267                );
1268            }
1269            other => panic!("Expected ErrorResponse, was {other:?}"),
1270        }
1271    }
1272
1273    #[rstest]
1274    fn test_parse_frame_order_response() {
1275        let value = load_value("ws_order_response.json");
1276        let frame = parse_bybit_ws_frame(value);
1277        match frame {
1278            BybitWsFrame::OrderResponse(resp) => {
1279                assert_eq!(resp.op, "order.create");
1280                assert_eq!(resp.ret_code, 0);
1281                assert_eq!(resp.ret_msg, "OK");
1282            }
1283            other => panic!("Expected OrderResponse, was {other:?}"),
1284        }
1285    }
1286
1287    #[rstest]
1288    fn test_parse_frame_orderbook() {
1289        let value = load_value("ws_orderbook_snapshot.json");
1290        let frame = parse_bybit_ws_frame(value);
1291        assert!(
1292            matches!(frame, BybitWsFrame::Orderbook(_)),
1293            "Expected Orderbook, was {frame:?}"
1294        );
1295    }
1296
1297    #[rstest]
1298    fn test_parse_frame_trade() {
1299        let value = load_value("ws_public_trade.json");
1300        let frame = parse_bybit_ws_frame(value);
1301        assert!(
1302            matches!(frame, BybitWsFrame::Trade(_)),
1303            "Expected Trade, was {frame:?}"
1304        );
1305    }
1306
1307    #[rstest]
1308    fn test_parse_frame_kline() {
1309        let value = load_value("ws_kline.json");
1310        let frame = parse_bybit_ws_frame(value);
1311        assert!(
1312            matches!(frame, BybitWsFrame::Kline(_)),
1313            "Expected Kline, was {frame:?}"
1314        );
1315    }
1316
1317    #[rstest]
1318    fn test_parse_frame_ticker_linear() {
1319        let value = load_value("ws_ticker_linear.json");
1320        let frame = parse_bybit_ws_frame(value);
1321        assert!(
1322            matches!(frame, BybitWsFrame::TickerLinear(_)),
1323            "Expected TickerLinear, was {frame:?}"
1324        );
1325    }
1326
1327    #[rstest]
1328    fn test_parse_frame_ticker_option() {
1329        let value = load_value("ws_ticker_option.json");
1330        let frame = parse_bybit_ws_frame(value);
1331        assert!(
1332            matches!(frame, BybitWsFrame::TickerOption(_)),
1333            "Expected TickerOption, was {frame:?}"
1334        );
1335    }
1336
1337    #[rstest]
1338    fn test_parse_frame_account_order() {
1339        let value = load_value("ws_account_order.json");
1340        let frame = parse_bybit_ws_frame(value);
1341        assert!(
1342            matches!(frame, BybitWsFrame::AccountOrder(_)),
1343            "Expected AccountOrder, was {frame:?}"
1344        );
1345    }
1346
1347    #[rstest]
1348    fn test_parse_frame_account_execution() {
1349        let value = load_value("ws_account_execution.json");
1350        let frame = parse_bybit_ws_frame(value);
1351        assert!(
1352            matches!(frame, BybitWsFrame::AccountExecution(_)),
1353            "Expected AccountExecution, was {frame:?}"
1354        );
1355    }
1356
1357    #[rstest]
1358    fn test_parse_frame_account_wallet() {
1359        let value = load_value("ws_account_wallet.json");
1360        let frame = parse_bybit_ws_frame(value);
1361        assert!(
1362            matches!(frame, BybitWsFrame::AccountWallet(_)),
1363            "Expected AccountWallet, was {frame:?}"
1364        );
1365    }
1366
1367    #[rstest]
1368    fn test_parse_frame_account_position() {
1369        let value = load_value("ws_account_position.json");
1370        let frame = parse_bybit_ws_frame(value);
1371        assert!(
1372            matches!(frame, BybitWsFrame::AccountPosition(_)),
1373            "Expected AccountPosition, was {frame:?}"
1374        );
1375    }
1376
1377    #[rstest]
1378    fn test_parse_frame_unknown_message() {
1379        let value: serde_json::Value = serde_json::json!({"foo": "bar"});
1380        let frame = parse_bybit_ws_frame(value);
1381        assert!(
1382            matches!(frame, BybitWsFrame::Unknown(_)),
1383            "Expected Unknown, was {frame:?}"
1384        );
1385    }
1386
1387    #[rstest]
1388    fn test_parse_raw_reconnected_signal() {
1389        let msg = Message::Text(nautilus_network::RECONNECTED.to_string().into());
1390        let result = BybitWsFeedHandler::parse_raw_frame(msg);
1391        assert!(
1392            matches!(result, Some(BybitWsFrame::Reconnected)),
1393            "Expected Some(Reconnected), was {result:?}"
1394        );
1395    }
1396
1397    #[rstest]
1398    fn test_parse_raw_pong_text() {
1399        let msg = Message::Text("pong".into());
1400        let result = BybitWsFeedHandler::parse_raw_frame(msg);
1401        assert!(result.is_none(), "Expected None for pong, was {result:?}");
1402    }
1403
1404    #[rstest]
1405    fn test_parse_raw_json_pong_message() {
1406        let msg = Message::Text(
1407            r#"{"args":["1777226678908"],"conn_id":"yzr7jz02gws1vh60mk5m-hxqdp","op":"pong"}"#
1408                .into(),
1409        );
1410        let result = BybitWsFeedHandler::parse_raw_frame(msg);
1411        assert!(
1412            result.is_none(),
1413            "Expected None for JSON pong, was {result:?}"
1414        );
1415    }
1416
1417    #[rstest]
1418    fn test_parse_raw_valid_json() {
1419        let json = load_test_json("ws_public_trade.json");
1420        let msg = Message::Text(json.into());
1421        let result = BybitWsFeedHandler::parse_raw_frame(msg);
1422        assert!(
1423            matches!(result, Some(BybitWsFrame::Trade(_))),
1424            "Expected Some(Trade), was {result:?}"
1425        );
1426    }
1427
1428    #[rstest]
1429    fn test_parse_raw_invalid_json() {
1430        let msg = Message::Text("not valid json".into());
1431        let result = BybitWsFeedHandler::parse_raw_frame(msg);
1432        assert!(
1433            result.is_none(),
1434            "Expected None for invalid JSON, was {result:?}"
1435        );
1436    }
1437
1438    #[rstest]
1439    fn test_parse_raw_binary_message() {
1440        let msg = Message::Binary(vec![0x01, 0x02].into());
1441        let result = BybitWsFeedHandler::parse_raw_frame(msg);
1442        assert!(result.is_none(), "Expected None for binary, was {result:?}");
1443    }
1444
1445    #[rstest]
1446    fn test_subscription_ack_with_req_id_confirms_only_that_topic() {
1447        let handler = create_test_handler();
1448        handler.subscriptions.mark_subscribe("orderbook.50.BTCUSDT");
1449        handler.subscriptions.mark_subscribe("publicTrade.BTCUSDT");
1450
1451        let ack = BybitWsSubscriptionMsg {
1452            success: true,
1453            op: BybitWsOperation::Subscribe,
1454            conn_id: None,
1455            req_id: Some("orderbook.50.BTCUSDT".to_string()),
1456            ret_msg: None,
1457        };
1458
1459        handler.handle_subscription_ack(&ack);
1460
1461        // Only orderbook should be confirmed, trade stays pending
1462        assert!(
1463            handler
1464                .subscriptions
1465                .pending_subscribe_topics()
1466                .contains(&"publicTrade.BTCUSDT".to_string())
1467        );
1468        assert!(
1469            !handler
1470                .subscriptions
1471                .pending_subscribe_topics()
1472                .contains(&"orderbook.50.BTCUSDT".to_string())
1473        );
1474    }
1475
1476    #[rstest]
1477    fn test_subscription_failure_with_req_id_marks_only_that_topic() {
1478        let handler = create_test_handler();
1479        handler.subscriptions.mark_subscribe("orderbook.50.BTCUSDT");
1480        handler.subscriptions.mark_subscribe("publicTrade.BTCUSDT");
1481
1482        let ack = BybitWsSubscriptionMsg {
1483            success: false,
1484            op: BybitWsOperation::Subscribe,
1485            conn_id: None,
1486            req_id: Some("orderbook.50.BTCUSDT".to_string()),
1487            ret_msg: Some("Invalid topic".to_string()),
1488        };
1489
1490        handler.handle_subscription_ack(&ack);
1491
1492        // Orderbook should be marked as failed (back to pending for retry)
1493        // Trade should remain pending (unaffected)
1494        let pending = handler.subscriptions.pending_subscribe_topics();
1495        assert!(pending.contains(&"orderbook.50.BTCUSDT".to_string()));
1496        assert!(pending.contains(&"publicTrade.BTCUSDT".to_string()));
1497    }
1498
1499    #[rstest]
1500    fn test_error_response_with_subscribe_op_triggers_mark_failure() {
1501        let handler = create_test_handler();
1502        handler
1503            .subscriptions
1504            .mark_subscribe("invalid.topic.BTCUSDT");
1505
1506        let resp = BybitWsResponse {
1507            op: Some(BybitWsOperation::Subscribe),
1508            topic: None,
1509            success: Some(false),
1510            conn_id: None,
1511            req_id: Some("invalid.topic.BTCUSDT".to_string()),
1512            ret_code: Some(10001),
1513            ret_msg: Some("Invalid topic".to_string()),
1514        };
1515
1516        handler.handle_subscription_error(&resp);
1517
1518        // Topic should still be in pending (mark_failure moves confirmed -> pending)
1519        let pending = handler.subscriptions.pending_subscribe_topics();
1520        assert!(pending.contains(&"invalid.topic.BTCUSDT".to_string()));
1521    }
1522
1523    #[rstest]
1524    fn test_already_subscribed_error_confirms_topic() {
1525        let handler = create_test_handler();
1526        handler.subscriptions.mark_subscribe("tickers.ETHUSDT");
1527
1528        let resp = BybitWsResponse {
1529            op: Some(BybitWsOperation::Subscribe),
1530            topic: None,
1531            success: Some(false),
1532            conn_id: None,
1533            req_id: Some("tickers.ETHUSDT".to_string()),
1534            ret_code: Some(10001),
1535            ret_msg: Some("error:already subscribed,topic:tickers.ETHUSDT".to_string()),
1536        };
1537
1538        handler.handle_subscription_error(&resp);
1539
1540        let pending = handler.subscriptions.pending_subscribe_topics();
1541        assert!(!pending.contains(&"tickers.ETHUSDT".to_string()));
1542        let symbols = handler.subscriptions.confirmed();
1543        let entry = symbols
1544            .get(&Ustr::from("tickers"))
1545            .expect("channel present");
1546        assert!(entry.contains(&Ustr::from("ETHUSDT")));
1547    }
1548
1549    #[rstest]
1550    fn test_subscription_ack_without_req_id_confirms_all_pending() {
1551        let handler = create_test_handler();
1552        handler.subscriptions.mark_subscribe("orderbook.50.BTCUSDT");
1553        handler.subscriptions.mark_subscribe("publicTrade.BTCUSDT");
1554
1555        let ack = BybitWsSubscriptionMsg {
1556            success: true,
1557            op: BybitWsOperation::Subscribe,
1558            conn_id: None,
1559            req_id: None,
1560            ret_msg: None,
1561        };
1562
1563        handler.handle_subscription_ack(&ack);
1564
1565        // Both should be confirmed when no req_id
1566        assert!(handler.subscriptions.pending_subscribe_topics().is_empty());
1567    }
1568}
nautilus_bybit/websocket/handler.rs

nautilus_bybit/websocket/
handler.rs
