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_network::{
28    error::SendError,
29    retry::{RetryManager, create_websocket_retry_manager},
30    websocket::{AuthTracker, SubscriptionState, WebSocketClient},
31};
32use serde::Serialize;
33use serde_json::Value;
34use tokio_tungstenite::tungstenite::Message;
35use ustr::Ustr;
36
37use super::{
38    enums::BybitWsOperation,
39    error::{BybitWsError, create_bybit_timeout_error, should_retry_bybit_error},
40    messages::{
41        BybitWebSocketError, BybitWsFrame, BybitWsMessage, BybitWsOrderResponse, BybitWsResponse,
42        BybitWsSubscriptionMsg,
43    },
44    parse::parse_bybit_ws_frame,
45};
46use crate::common::{
47    enums::{BybitProductType, BybitWsOrderRequestOp},
48    rate_limit::{
49        BYBIT_RATE_LIMIT_HEADER, BYBIT_RATE_LIMIT_RESET_HEADER, BYBIT_RATE_LIMIT_STATUS_HEADER,
50        BybitRateLimiter,
51    },
52};
53
54/// Semantic order command whose time-sensitive header is built at send time.
55#[derive(Debug)]
56pub struct BybitWsOrderCommand {
57    pub(crate) req_id: String,
58    pub(crate) op: BybitWsOrderRequestOp,
59    pub(crate) category: BybitProductType,
60    pub(crate) weight: u32,
61    pub(crate) referer: Option<String>,
62    pub(crate) args: Vec<Value>,
63}
64
65#[derive(Debug, Serialize)]
66#[serde(rename_all = "camelCase")]
67struct BybitWsSignedOrderRequest {
68    req_id: String,
69    op: BybitWsOrderRequestOp,
70    header: BybitWsSignedOrderHeader,
71    args: Vec<Value>,
72}
73
74#[derive(Debug, Serialize)]
75#[serde(rename_all = "SCREAMING-KEBAB-CASE")]
76struct BybitWsSignedOrderHeader {
77    x_bapi_timestamp: String,
78    x_bapi_recv_window: String,
79    #[serde(rename = "Referer", skip_serializing_if = "Option::is_none")]
80    referer: Option<String>,
81}
82
83#[derive(Clone, Copy, Debug)]
84struct PendingRate {
85    endpoint: &'static str,
86    category: BybitProductType,
87}
88
89#[derive(Debug)]
90enum OrderSendFailure {
91    NotSent(BybitWsError),
92    Ambiguous(BybitWsError),
93}
94
95/// Commands sent from the outer client to the inner message handler.
96#[derive(Debug)]
97pub enum HandlerCommand {
98    SetClient(WebSocketClient),
99    Disconnect,
100    Authenticate { payload: String },
101    Subscribe { topics: Vec<String> },
102    Unsubscribe { topics: Vec<String> },
103    SendOrder { command: BybitWsOrderCommand },
104    SendOrders { commands: Vec<BybitWsOrderCommand> },
105}
106
107pub(super) struct BybitWsFeedHandler {
108    signal: Arc<AtomicBool>,
109    inner: Option<WebSocketClient>,
110    cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
111    raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
112    auth_tracker: AuthTracker,
113    subscriptions: SubscriptionState,
114    rate_limiter: BybitRateLimiter,
115    recv_window_ms: Arc<AtomicU64>,
116    pending_rates: DashMap<String, PendingRate>,
117    pending_orders: VecDeque<BybitWsOrderCommand>,
118    retry_manager: RetryManager<BybitWsError>,
119}
120
121impl BybitWsFeedHandler {
122    /// Creates a new [`BybitWsFeedHandler`] instance.
123    pub(super) fn new(
124        signal: Arc<AtomicBool>,
125        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
126        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
127        auth_tracker: AuthTracker,
128        subscriptions: SubscriptionState,
129        rate_limiter: BybitRateLimiter,
130        recv_window_ms: Arc<AtomicU64>,
131    ) -> Self {
132        Self {
133            signal,
134            inner: None,
135            cmd_rx,
136            raw_rx,
137            auth_tracker,
138            subscriptions,
139            rate_limiter,
140            recv_window_ms,
141            pending_rates: DashMap::new(),
142            pending_orders: VecDeque::new(),
143            retry_manager: create_websocket_retry_manager(),
144        }
145    }
146
147    pub(super) fn is_stopped(&self) -> bool {
148        self.signal.load(Ordering::Relaxed)
149    }
150
151    /// Sends a WebSocket message with retry logic.
152    async fn send_with_retry(&self, payload: String) -> Result<(), BybitWsError> {
153        if let Some(client) = &self.inner {
154            self.retry_manager
155                .execute_with_retry(
156                    "websocket_send",
157                    || {
158                        let payload = payload.clone();
159                        async move {
160                            client
161                                .send_text(payload, None)
162                                .await
163                                .map_err(|e| BybitWsError::Transport(format!("Send failed: {e}")))
164                        }
165                    },
166                    should_retry_bybit_error,
167                    |e| create_bybit_timeout_error(e.to_string()),
168                )
169                .await
170        } else {
171            Err(BybitWsError::ClientError(
172                "No active WebSocket client".to_string(),
173            ))
174        }
175    }
176
177    async fn send_order(&self, command: &BybitWsOrderCommand) -> Result<bool, OrderSendFailure> {
178        if !self.auth_tracker.is_authenticated() {
179            return Ok(false);
180        }
181
182        let client = self.inner.as_ref().ok_or_else(|| {
183            OrderSendFailure::NotSent(BybitWsError::ClientError(
184                "No active WebSocket client".to_string(),
185            ))
186        })?;
187        let (endpoint, _) = order_operation(command.op);
188
189        self.rate_limiter
190            .acquire_ws_order(endpoint, command.category, command.weight)
191            .await
192            .map_err(|e| OrderSendFailure::NotSent(BybitWsError::ClientError(e)))?;
193
194        if !self.auth_tracker.is_authenticated() {
195            return Ok(false);
196        }
197
198        let connection_epoch = client.connection_epoch();
199        let recv_window_ms = self.recv_window_ms.load(Ordering::Acquire);
200        let request = BybitWsSignedOrderRequest {
201            req_id: command.req_id.clone(),
202            op: command.op,
203            header: BybitWsSignedOrderHeader {
204                x_bapi_timestamp: nautilus_core::time::get_atomic_clock_realtime()
205                    .get_time_ms()
206                    .to_string(),
207                x_bapi_recv_window: recv_window_ms.to_string(),
208                referer: command.referer.clone(),
209            },
210            args: command.args.clone(),
211        };
212        let payload = serde_json::to_string(&request)
213            .map_err(|e| OrderSendFailure::NotSent(BybitWsError::Json(e.to_string())))?;
214        self.pending_rates.insert(
215            command.req_id.clone(),
216            PendingRate {
217                endpoint,
218                category: command.category,
219            },
220        );
221
222        match client
223            .send_text_on_connection(payload, None, connection_epoch)
224            .await
225        {
226            Ok(()) => Ok(true),
227            Err(e) => self.classify_order_send_error(&command.req_id, e),
228        }
229    }
230
231    fn classify_order_send_error(
232        &self,
233        req_id: &str,
234        error: SendError,
235    ) -> Result<bool, OrderSendFailure> {
236        match error {
237            SendError::ConnectionChanged => {
238                self.pending_rates.remove(req_id);
239                self.auth_tracker.invalidate();
240                Ok(false)
241            }
242            SendError::Timeout => {
243                self.pending_rates.remove(req_id);
244                self.auth_tracker.invalidate();
245                Err(OrderSendFailure::NotSent(BybitWsError::ClientError(
246                    "Order command was not written".to_string(),
247                )))
248            }
249            SendError::InvalidInput(_) | SendError::Closed => {
250                self.pending_rates.remove(req_id);
251                Err(OrderSendFailure::NotSent(BybitWsError::ClientError(
252                    "Order command was not written".to_string(),
253                )))
254            }
255            error => Err(OrderSendFailure::Ambiguous(BybitWsError::Send(
256                error.to_string(),
257            ))),
258        }
259    }
260
261    pub(super) async fn next(&mut self) -> Option<BybitWsMessage> {
262        loop {
263            if self.auth_tracker.is_authenticated()
264                && let Some(command) = self.pending_orders.pop_front()
265            {
266                let req_id = command.req_id.clone();
267                match self.send_order(&command).await {
268                    Ok(true) => {}
269                    Ok(false) => self.pending_orders.push_front(command),
270                    Err(OrderSendFailure::NotSent(error)) => {
271                        return Some(order_not_sent_message(&command, &error));
272                    }
273                    Err(OrderSendFailure::Ambiguous(error)) => {
274                        log::error!("Ambiguous order send failure: req_id={req_id}, error={error}");
275                    }
276                }
277                continue;
278            }
279
280            tokio::select! {
281                Some(cmd) = self.cmd_rx.recv() => {
282                    match cmd {
283                        HandlerCommand::SetClient(client) => {
284                            log::debug!("WebSocketClient received by handler");
285                            self.inner = Some(client);
286                        }
287                        HandlerCommand::Disconnect => {
288                            log::debug!("Disconnect command received");
289
290                            if let Some(client) = self.inner.take() {
291                                client.disconnect().await;
292                            }
293                        }
294                        HandlerCommand::Authenticate { payload } => {
295                            log::debug!("Authenticate command received");
296
297                            if let Err(e) = self.send_with_retry(payload).await {
298                                log::error!("Failed to send authentication after retries: {e}");
299                            }
300                        }
301                        HandlerCommand::Subscribe { topics } => {
302                            for topic in topics {
303                                log::debug!("Subscribing to topic: topic={topic}");
304                                if let Err(e) = self.send_with_retry(topic.clone()).await {
305                                    log::error!("Failed to send subscription after retries: topic={topic}, error={e}");
306                                }
307                            }
308                        }
309                        HandlerCommand::Unsubscribe { topics } => {
310                            for topic in topics {
311                                log::debug!("Unsubscribing from topic: topic={topic}");
312                                if let Err(e) = self.send_with_retry(topic.clone()).await {
313                                    log::error!("Failed to send unsubscription after retries: topic={topic}, error={e}");
314                                }
315                            }
316                        }
317                        HandlerCommand::SendOrder { command } => {
318                            let req_id = command.req_id.clone();
319                            match self.send_order(&command).await {
320                                Ok(true) => {}
321                                Ok(false) => self.pending_orders.push_back(command),
322                                Err(OrderSendFailure::NotSent(error)) => {
323                                    return Some(order_not_sent_message(&command, &error));
324                                }
325                                Err(OrderSendFailure::Ambiguous(error)) => {
326                                    log::error!(
327                                        "Ambiguous order send failure: req_id={req_id}, error={error}"
328                                    );
329                                }
330                            }
331                        }
332                        HandlerCommand::SendOrders { commands } => {
333                            self.pending_orders.extend(commands);
334                        }
335                    }
336                }
337
338                () = tokio::time::sleep(std::time::Duration::from_millis(100)) => {
339                    if self.signal.load(Ordering::Relaxed) {
340                        log::debug!("Stop signal received during idle period");
341                        return None;
342                    }
343                }
344
345                msg = self.raw_rx.recv() => {
346                    let msg = match msg {
347                        Some(msg) => msg,
348                        None => {
349                            log::debug!("WebSocket stream closed");
350                            return None;
351                        }
352                    };
353
354                    if let Message::Ping(data) = &msg {
355                        log::trace!("Received ping frame with {} bytes", data.len());
356
357                        if let Some(client) = &self.inner
358                            && let Err(e) = client.send_pong(data.to_vec()).await
359                        {
360                            log::warn!("Failed to send pong frame: error={e}");
361                        }
362                        continue;
363                    }
364
365                    let frame = match Self::parse_raw_frame(msg) {
366                        Some(frame) => frame,
367                        None => continue,
368                    };
369
370                    if self.signal.load(Ordering::Relaxed) {
371                        log::debug!("Stop signal received");
372                        return None;
373                    }
374
375                    match frame {
376                        BybitWsFrame::Subscription(ref sub_msg) => {
377                            self.handle_subscription_ack(sub_msg);
378                        }
379                        BybitWsFrame::Auth(auth_response) => {
380                            let is_success = auth_response.success.unwrap_or(false)
381                                || (auth_response.ret_code == Some(0));
382
383                            if is_success {
384                                self.auth_tracker.succeed();
385                                log::debug!("WebSocket authenticated");
386                            } else {
387                                let error_msg = auth_response
388                                    .ret_msg
389                                    .as_deref()
390                                    .unwrap_or("Authentication rejected");
391                                self.auth_tracker.fail(error_msg);
392                                log::error!("WebSocket authentication failed: error={error_msg}");
393                            }
394                            return Some(BybitWsMessage::Auth(auth_response));
395                        }
396                        BybitWsFrame::ErrorResponse(ref resp) => {
397                            // Failed subscription/unsubscription ACKs arrive as
398                            // ErrorResponse when success=false. Route them through
399                            // the subscription state machine when the op field is
400                            // present, otherwise forward as a generic error.
401                            if let Some(op) = &resp.op {
402                                if *op == BybitWsOperation::Subscribe
403                                    || *op == BybitWsOperation::Unsubscribe
404                                {
405                                    self.handle_subscription_error(resp);
406                                } else {
407                                    let error = BybitWebSocketError::from_response(resp);
408                                    return Some(BybitWsMessage::Error(error));
409                                }
410                            } else {
411                                let error = BybitWebSocketError::from_response(resp);
412                                return Some(BybitWsMessage::Error(error));
413                            }
414                        }
415                        BybitWsFrame::OrderResponse(resp) => {
416                            self.observe_order_rate_limit(&resp);
417                            return Some(BybitWsMessage::OrderResponse(resp));
418                        }
419                        BybitWsFrame::Orderbook(msg) => {
420                            return Some(BybitWsMessage::Orderbook(msg));
421                        }
422                        BybitWsFrame::Trade(msg) => {
423                            return Some(BybitWsMessage::Trade(msg));
424                        }
425                        BybitWsFrame::Kline(msg) => {
426                            return Some(BybitWsMessage::Kline(msg));
427                        }
428                        BybitWsFrame::TickerLinear(msg) => {
429                            return Some(BybitWsMessage::TickerLinear(msg));
430                        }
431                        BybitWsFrame::TickerOption(msg) => {
432                            return Some(BybitWsMessage::TickerOption(msg));
433                        }
434                        BybitWsFrame::AccountOrder(msg) => {
435                            return Some(BybitWsMessage::AccountOrder(msg));
436                        }
437                        BybitWsFrame::AccountExecution(msg) => {
438                            return Some(BybitWsMessage::AccountExecution(msg));
439                        }
440                        BybitWsFrame::AccountExecutionFast(msg) => {
441                            return Some(BybitWsMessage::AccountExecutionFast(msg));
442                        }
443                        BybitWsFrame::AccountWallet(msg) => {
444                            return Some(BybitWsMessage::AccountWallet(msg));
445                        }
446                        BybitWsFrame::AccountPosition(msg) => {
447                            return Some(BybitWsMessage::AccountPosition(msg));
448                        }
449                        BybitWsFrame::Reconnected => {
450                            self.auth_tracker.invalidate();
451                            return Some(BybitWsMessage::Reconnected);
452                        }
453                        BybitWsFrame::Unknown(value) => {
454                            log::debug!("Unknown WebSocket frame: {value}");
455                        }
456                    }
457                }
458            }
459        }
460    }
461
462    fn observe_order_rate_limit(&self, response: &super::messages::BybitWsOrderResponse) {
463        let Some(req_id) = response.req_id.as_deref() else {
464            return;
465        };
466        let Some((_, pending)) = self.pending_rates.remove(req_id) else {
467            return;
468        };
469        let Some(header) = response.header.as_ref() else {
470            return;
471        };
472        let parse_u32 = |key: &str| {
473            header.get(key).and_then(|value| {
474                value
475                    .as_u64()
476                    .and_then(|value| u32::try_from(value).ok())
477                    .or_else(|| value.as_str()?.parse::<u32>().ok())
478            })
479        };
480        let Some(limit) = parse_u32(BYBIT_RATE_LIMIT_HEADER) else {
481            return;
482        };
483        let Some(remaining) = parse_u32(BYBIT_RATE_LIMIT_STATUS_HEADER) else {
484            return;
485        };
486        let reset_timestamp_ms = header.get(BYBIT_RATE_LIMIT_RESET_HEADER).and_then(|value| {
487            value
488                .as_i64()
489                .or_else(|| value.as_str()?.parse::<i64>().ok())
490        });
491        self.rate_limiter.observe_account(
492            pending.endpoint,
493            Some(pending.category),
494            limit,
495            remaining,
496            reset_timestamp_ms,
497        );
498    }
499
500    fn handle_subscription_ack(&self, sub_msg: &BybitWsSubscriptionMsg) {
501        match sub_msg.op {
502            BybitWsOperation::Subscribe => {
503                if sub_msg.success {
504                    if let Some(topic) = &sub_msg.req_id {
505                        self.subscriptions.confirm_subscribe(topic);
506                        log::debug!("Subscription confirmed: topic={topic}");
507                    } else {
508                        // No req_id, fall back to confirming all pending
509                        for topic in self.subscriptions.pending_subscribe_topics() {
510                            self.subscriptions.confirm_subscribe(&topic);
511                            log::debug!("Subscription confirmed (bulk): topic={topic}");
512                        }
513                    }
514                } else if let Some(topic) = &sub_msg.req_id {
515                    self.subscriptions.mark_failure(topic);
516                    log::warn!(
517                        "Subscription failed: topic={topic}, error={:?}",
518                        sub_msg.ret_msg
519                    );
520                } else {
521                    for topic in self.subscriptions.pending_subscribe_topics() {
522                        self.subscriptions.mark_failure(&topic);
523                        log::warn!(
524                            "Subscription failed (bulk): topic={topic}, error={:?}",
525                            sub_msg.ret_msg
526                        );
527                    }
528                }
529            }
530            BybitWsOperation::Unsubscribe => {
531                if sub_msg.success {
532                    if let Some(topic) = &sub_msg.req_id {
533                        self.subscriptions.confirm_unsubscribe(topic);
534                        log::debug!("Unsubscription confirmed: topic={topic}");
535                    } else {
536                        for topic in self.subscriptions.pending_unsubscribe_topics() {
537                            self.subscriptions.confirm_unsubscribe(&topic);
538                            log::debug!("Unsubscription confirmed (bulk): topic={topic}");
539                        }
540                    }
541                } else {
542                    let topic_desc = sub_msg.req_id.as_deref().unwrap_or("unknown");
543                    log::warn!(
544                        "Unsubscription failed: topic={topic_desc}, error={:?}",
545                        sub_msg.ret_msg
546                    );
547                }
548            }
549            _ => {}
550        }
551    }
552
553    fn handle_subscription_error(&self, resp: &BybitWsResponse) {
554        let topic = resp.req_id.as_deref().unwrap_or("unknown");
555        let error_msg = resp.ret_msg.as_deref().unwrap_or("unknown error");
556
557        match resp.op {
558            Some(BybitWsOperation::Subscribe) => {
559                // Duplicate subscribe is harmless: the topic is active on the
560                // venue, so confirm it instead of looping retries every reconnect.
561                if is_already_subscribed_error(error_msg)
562                    && let Some(ref req_id) = resp.req_id
563                {
564                    self.subscriptions.confirm_subscribe(req_id);
565                    log::debug!("Subscription duplicate ignored: topic={topic}, error={error_msg}");
566                    return;
567                }
568
569                if let Some(ref req_id) = resp.req_id {
570                    self.subscriptions.mark_failure(req_id);
571                } else {
572                    for t in self.subscriptions.pending_subscribe_topics() {
573                        self.subscriptions.mark_failure(&t);
574                    }
575                }
576                log::warn!("Subscription error: topic={topic}, error={error_msg}");
577            }
578            Some(BybitWsOperation::Unsubscribe) => {
579                log::warn!("Unsubscription error: topic={topic}, error={error_msg}");
580            }
581            _ => {}
582        }
583    }
584
585    fn parse_raw_frame(msg: Message) -> Option<BybitWsFrame> {
586        match msg {
587            Message::Text(text) => {
588                if text == nautilus_network::RECONNECTED {
589                    log::debug!("Received WebSocket reconnected signal");
590                    return Some(BybitWsFrame::Reconnected);
591                }
592
593                if text.trim().eq_ignore_ascii_case("pong") {
594                    return None;
595                }
596
597                log::trace!("Raw websocket message: {text}");
598
599                let value: serde_json::Value = match serde_json::from_str(&text) {
600                    Ok(v) => v,
601                    Err(e) => {
602                        log::error!("Failed to parse WebSocket message: {e}: {text}");
603                        return None;
604                    }
605                };
606
607                if value
608                    .get("op")
609                    .and_then(serde_json::Value::as_str)
610                    .is_some_and(|op| op == BybitWsOperation::Pong.as_ref())
611                {
612                    return None;
613                }
614
615                Some(parse_bybit_ws_frame(value))
616            }
617            Message::Binary(msg) => {
618                log::debug!("Raw binary frame ({} bytes)", msg.len());
619                log::trace!("Raw binary: {msg:?}");
620                None
621            }
622            Message::Close(_) => {
623                log::debug!("Received close message, waiting for reconnection");
624                None
625            }
626            _ => None,
627        }
628    }
629}
630
631fn order_operation(op: BybitWsOrderRequestOp) -> (&'static str, &'static str) {
632    match op {
633        BybitWsOrderRequestOp::Create => ("/v5/order/create", "order.create"),
634        BybitWsOrderRequestOp::Amend => ("/v5/order/amend", "order.amend"),
635        BybitWsOrderRequestOp::Cancel => ("/v5/order/cancel", "order.cancel"),
636        BybitWsOrderRequestOp::CreateBatch => ("/v5/order/create-batch", "order.create-batch"),
637        BybitWsOrderRequestOp::AmendBatch => ("/v5/order/amend-batch", "order.amend-batch"),
638        BybitWsOrderRequestOp::CancelBatch => ("/v5/order/cancel-batch", "order.cancel-batch"),
639    }
640}
641
642fn order_not_sent_message(command: &BybitWsOrderCommand, error: &BybitWsError) -> BybitWsMessage {
643    let (_, op) = order_operation(command.op);
644    BybitWsMessage::OrderResponse(BybitWsOrderResponse {
645        op: Ustr::from(op),
646        conn_id: None,
647        ret_code: -1,
648        ret_msg: error.to_string(),
649        data: Value::Object(serde_json::Map::new()),
650        req_id: Some(command.req_id.clone()),
651        header: None,
652        ret_ext_info: None,
653    })
654}
655
656fn is_already_subscribed_error(error_msg: &str) -> bool {
657    error_msg
658        .to_ascii_lowercase()
659        .contains("already subscribed")
660}
661
662#[cfg(test)]
663mod tests {
664    use rstest::rstest;
665    use ustr::Ustr;
666
667    use super::*;
668    use crate::common::{
669        consts::BYBIT_WS_TOPIC_DELIMITER, rate_limit::BybitRateLimiter, testing::load_test_json,
670    };
671
672    fn create_test_handler() -> BybitWsFeedHandler {
673        let signal = Arc::new(AtomicBool::new(false));
674        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
675        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
676        let auth_tracker = AuthTracker::new();
677        let subscriptions = SubscriptionState::new(BYBIT_WS_TOPIC_DELIMITER);
678
679        BybitWsFeedHandler::new(
680            signal,
681            cmd_rx,
682            raw_rx,
683            auth_tracker,
684            subscriptions,
685            BybitRateLimiter::for_websocket(
686                "wss://bybit-handler-test.invalid",
687                Some("test-key"),
688                None,
689            ),
690            Arc::new(AtomicU64::new(5_000)),
691        )
692    }
693
694    fn load_value(fixture: &str) -> serde_json::Value {
695        let json = load_test_json(fixture);
696        serde_json::from_str(&json).unwrap()
697    }
698
699    #[rstest]
700    fn test_handler_initializes() {
701        let _handler = create_test_handler();
702    }
703
704    #[tokio::test]
705    async fn order_without_writer_produces_correlated_terminal_response() {
706        let handler = create_test_handler();
707        handler.auth_tracker.succeed();
708        let command = BybitWsOrderCommand {
709            req_id: "not-sent-request".to_string(),
710            op: BybitWsOrderRequestOp::Create,
711            category: BybitProductType::Linear,
712            weight: 1,
713            referer: None,
714            args: vec![serde_json::json!({"category": "linear"})],
715        };
716
717        let failure = handler.send_order(&command).await.unwrap_err();
718        let OrderSendFailure::NotSent(error) = failure else {
719            panic!("Expected definitive not-sent failure, was {failure:?}");
720        };
721        let BybitWsMessage::OrderResponse(response) = order_not_sent_message(&command, &error)
722        else {
723            panic!("Expected order response");
724        };
725
726        assert_eq!(response.op.as_str(), "order.create");
727        assert_eq!(response.ret_code, -1);
728        assert_eq!(response.req_id.as_deref(), Some("not-sent-request"));
729        assert_eq!(response.ret_msg, "Client error: No active WebSocket client");
730    }
731
732    #[rstest]
733    fn connection_change_invalidates_authentication_before_requeue() {
734        let handler = create_test_handler();
735        handler.auth_tracker.succeed();
736        handler.pending_rates.insert(
737            "reconnect-request".to_string(),
738            PendingRate {
739                endpoint: "/v5/order/create",
740                category: BybitProductType::Linear,
741            },
742        );
743
744        let result = handler
745            .classify_order_send_error("reconnect-request", SendError::ConnectionChanged)
746            .unwrap();
747
748        assert!(!result);
749        assert!(!handler.auth_tracker.is_authenticated());
750        assert!(!handler.pending_rates.contains_key("reconnect-request"));
751    }
752
753    #[rstest]
754    fn pre_active_timeout_is_terminal_not_sent() {
755        let handler = create_test_handler();
756        handler.auth_tracker.succeed();
757        handler.pending_rates.insert(
758            "timeout-request".to_string(),
759            PendingRate {
760                endpoint: "/v5/order/create",
761                category: BybitProductType::Linear,
762            },
763        );
764
765        let failure = handler
766            .classify_order_send_error("timeout-request", SendError::Timeout)
767            .unwrap_err();
768
769        assert!(matches!(failure, OrderSendFailure::NotSent(_)));
770        assert!(!handler.auth_tracker.is_authenticated());
771        assert!(!handler.pending_rates.contains_key("timeout-request"));
772    }
773
774    #[rstest]
775    fn test_parse_frame_auth_success() {
776        let value = load_value("ws_auth_success.json");
777        let frame = parse_bybit_ws_frame(value);
778        match frame {
779            BybitWsFrame::Auth(auth) => {
780                assert_eq!(auth.conn_id.as_deref(), Some("cejreaspqfm9se7usbrg-2xh"));
781                assert_eq!(auth.ret_code, Some(0));
782                assert_eq!(auth.success, Some(true));
783            }
784            other => panic!("Expected Auth, was {other:?}"),
785        }
786    }
787
788    #[rstest]
789    fn test_parse_frame_auth_failure() {
790        let value = load_value("ws_auth_failure.json");
791        let frame = parse_bybit_ws_frame(value);
792        match frame {
793            BybitWsFrame::ErrorResponse(resp) => {
794                assert_eq!(resp.ret_code, Some(10003));
795                assert_eq!(resp.ret_msg.as_deref(), Some("Invalid apikey"));
796            }
797            other => panic!("Expected ErrorResponse, was {other:?}"),
798        }
799    }
800
801    #[rstest]
802    fn test_parse_frame_subscription_ack() {
803        let value = load_value("ws_subscription_ack.json");
804        let frame = parse_bybit_ws_frame(value);
805        match frame {
806            BybitWsFrame::Subscription(sub) => {
807                assert!(sub.success);
808                assert_eq!(sub.op, BybitWsOperation::Subscribe);
809                assert_eq!(sub.req_id.as_deref(), Some("sub-orderbook-1"));
810            }
811            other => panic!("Expected Subscription, was {other:?}"),
812        }
813    }
814
815    #[rstest]
816    fn test_parse_frame_subscription_failure() {
817        let value = load_value("ws_subscription_failure.json");
818        let frame = parse_bybit_ws_frame(value);
819        match frame {
820            BybitWsFrame::ErrorResponse(resp) => {
821                assert_eq!(
822                    resp.ret_msg.as_deref(),
823                    Some("Invalid topic: invalid.topic.BTCUSDT")
824                );
825            }
826            other => panic!("Expected ErrorResponse, was {other:?}"),
827        }
828    }
829
830    #[rstest]
831    fn test_parse_frame_order_response() {
832        let value = load_value("ws_order_response.json");
833        let frame = parse_bybit_ws_frame(value);
834        match frame {
835            BybitWsFrame::OrderResponse(resp) => {
836                assert_eq!(resp.op.as_str(), "order.create");
837                assert_eq!(resp.ret_code, 0);
838                assert_eq!(resp.ret_msg, "OK");
839            }
840            other => panic!("Expected OrderResponse, was {other:?}"),
841        }
842    }
843
844    #[rstest]
845    fn test_parse_frame_orderbook() {
846        let value = load_value("ws_orderbook_snapshot.json");
847        let frame = parse_bybit_ws_frame(value);
848        assert!(
849            matches!(frame, BybitWsFrame::Orderbook(_)),
850            "Expected Orderbook, was {frame:?}"
851        );
852    }
853
854    #[rstest]
855    fn test_parse_frame_trade() {
856        let value = load_value("ws_public_trade.json");
857        let frame = parse_bybit_ws_frame(value);
858        assert!(
859            matches!(frame, BybitWsFrame::Trade(_)),
860            "Expected Trade, was {frame:?}"
861        );
862    }
863
864    #[rstest]
865    fn test_parse_frame_kline() {
866        let value = load_value("ws_kline.json");
867        let frame = parse_bybit_ws_frame(value);
868        assert!(
869            matches!(frame, BybitWsFrame::Kline(_)),
870            "Expected Kline, was {frame:?}"
871        );
872    }
873
874    #[rstest]
875    fn test_parse_frame_ticker_linear() {
876        let value = load_value("ws_ticker_linear.json");
877        let frame = parse_bybit_ws_frame(value);
878        assert!(
879            matches!(frame, BybitWsFrame::TickerLinear(_)),
880            "Expected TickerLinear, was {frame:?}"
881        );
882    }
883
884    #[rstest]
885    fn test_parse_frame_ticker_option() {
886        let value = load_value("ws_ticker_option.json");
887        let frame = parse_bybit_ws_frame(value);
888        assert!(
889            matches!(frame, BybitWsFrame::TickerOption(_)),
890            "Expected TickerOption, was {frame:?}"
891        );
892    }
893
894    #[rstest]
895    fn test_parse_frame_account_order() {
896        let value = load_value("ws_account_order.json");
897        let frame = parse_bybit_ws_frame(value);
898        assert!(
899            matches!(frame, BybitWsFrame::AccountOrder(_)),
900            "Expected AccountOrder, was {frame:?}"
901        );
902    }
903
904    #[rstest]
905    fn test_parse_frame_account_execution() {
906        let value = load_value("ws_account_execution.json");
907        let frame = parse_bybit_ws_frame(value);
908        assert!(
909            matches!(frame, BybitWsFrame::AccountExecution(_)),
910            "Expected AccountExecution, was {frame:?}"
911        );
912    }
913
914    #[rstest]
915    fn test_parse_frame_account_wallet() {
916        let value = load_value("ws_account_wallet.json");
917        let frame = parse_bybit_ws_frame(value);
918        assert!(
919            matches!(frame, BybitWsFrame::AccountWallet(_)),
920            "Expected AccountWallet, was {frame:?}"
921        );
922    }
923
924    #[rstest]
925    fn test_parse_frame_account_position() {
926        let value = load_value("ws_account_position.json");
927        let frame = parse_bybit_ws_frame(value);
928        assert!(
929            matches!(frame, BybitWsFrame::AccountPosition(_)),
930            "Expected AccountPosition, was {frame:?}"
931        );
932    }
933
934    #[rstest]
935    fn test_parse_frame_unknown_message() {
936        let value: serde_json::Value = serde_json::json!({"foo": "bar"});
937        let frame = parse_bybit_ws_frame(value);
938        assert!(
939            matches!(frame, BybitWsFrame::Unknown(_)),
940            "Expected Unknown, was {frame:?}"
941        );
942    }
943
944    #[rstest]
945    fn test_parse_raw_reconnected_signal() {
946        let msg = Message::Text(nautilus_network::RECONNECTED.to_string().into());
947        let result = BybitWsFeedHandler::parse_raw_frame(msg);
948        assert!(
949            matches!(result, Some(BybitWsFrame::Reconnected)),
950            "Expected Some(Reconnected), was {result:?}"
951        );
952    }
953
954    #[rstest]
955    fn test_parse_raw_pong_text() {
956        let msg = Message::Text("pong".into());
957        let result = BybitWsFeedHandler::parse_raw_frame(msg);
958        assert!(result.is_none(), "Expected None for pong, was {result:?}");
959    }
960
961    #[rstest]
962    fn test_parse_raw_json_pong_message() {
963        let msg = Message::Text(
964            r#"{"args":["1777226678908"],"conn_id":"yzr7jz02gws1vh60mk5m-hxqdp","op":"pong"}"#
965                .into(),
966        );
967        let result = BybitWsFeedHandler::parse_raw_frame(msg);
968        assert!(
969            result.is_none(),
970            "Expected None for JSON pong, was {result:?}"
971        );
972    }
973
974    #[rstest]
975    fn test_parse_raw_valid_json() {
976        let json = load_test_json("ws_public_trade.json");
977        let msg = Message::Text(json.into());
978        let result = BybitWsFeedHandler::parse_raw_frame(msg);
979        assert!(
980            matches!(result, Some(BybitWsFrame::Trade(_))),
981            "Expected Some(Trade), was {result:?}"
982        );
983    }
984
985    #[rstest]
986    fn test_parse_raw_invalid_json() {
987        let msg = Message::Text("not valid json".into());
988        let result = BybitWsFeedHandler::parse_raw_frame(msg);
989        assert!(
990            result.is_none(),
991            "Expected None for invalid JSON, was {result:?}"
992        );
993    }
994
995    #[rstest]
996    fn test_parse_raw_binary_message() {
997        let msg = Message::Binary(vec![0x01, 0x02].into());
998        let result = BybitWsFeedHandler::parse_raw_frame(msg);
999        assert!(result.is_none(), "Expected None for binary, was {result:?}");
1000    }
1001
1002    #[rstest]
1003    fn test_subscription_ack_with_req_id_confirms_only_that_topic() {
1004        let handler = create_test_handler();
1005        handler.subscriptions.mark_subscribe("orderbook.50.BTCUSDT");
1006        handler.subscriptions.mark_subscribe("publicTrade.BTCUSDT");
1007
1008        let ack = BybitWsSubscriptionMsg {
1009            success: true,
1010            op: BybitWsOperation::Subscribe,
1011            conn_id: None,
1012            req_id: Some("orderbook.50.BTCUSDT".to_string()),
1013            ret_msg: None,
1014        };
1015
1016        handler.handle_subscription_ack(&ack);
1017
1018        // Only orderbook should be confirmed, trade stays pending
1019        assert!(
1020            handler
1021                .subscriptions
1022                .pending_subscribe_topics()
1023                .contains(&"publicTrade.BTCUSDT".to_string())
1024        );
1025        assert!(
1026            !handler
1027                .subscriptions
1028                .pending_subscribe_topics()
1029                .contains(&"orderbook.50.BTCUSDT".to_string())
1030        );
1031    }
1032
1033    #[rstest]
1034    fn test_subscription_failure_with_req_id_marks_only_that_topic() {
1035        let handler = create_test_handler();
1036        handler.subscriptions.mark_subscribe("orderbook.50.BTCUSDT");
1037        handler.subscriptions.mark_subscribe("publicTrade.BTCUSDT");
1038
1039        let ack = BybitWsSubscriptionMsg {
1040            success: false,
1041            op: BybitWsOperation::Subscribe,
1042            conn_id: None,
1043            req_id: Some("orderbook.50.BTCUSDT".to_string()),
1044            ret_msg: Some("Invalid topic".to_string()),
1045        };
1046
1047        handler.handle_subscription_ack(&ack);
1048
1049        // Orderbook should be marked as failed (back to pending for retry)
1050        // Trade should remain pending (unaffected)
1051        let pending = handler.subscriptions.pending_subscribe_topics();
1052        assert!(pending.contains(&"orderbook.50.BTCUSDT".to_string()));
1053        assert!(pending.contains(&"publicTrade.BTCUSDT".to_string()));
1054    }
1055
1056    #[rstest]
1057    fn test_error_response_with_subscribe_op_triggers_mark_failure() {
1058        let handler = create_test_handler();
1059        handler
1060            .subscriptions
1061            .mark_subscribe("invalid.topic.BTCUSDT");
1062
1063        let resp = BybitWsResponse {
1064            op: Some(BybitWsOperation::Subscribe),
1065            topic: None,
1066            success: Some(false),
1067            conn_id: None,
1068            req_id: Some("invalid.topic.BTCUSDT".to_string()),
1069            ret_code: Some(10001),
1070            ret_msg: Some("Invalid topic".to_string()),
1071        };
1072
1073        handler.handle_subscription_error(&resp);
1074
1075        // Topic should still be in pending (mark_failure moves confirmed -> pending)
1076        let pending = handler.subscriptions.pending_subscribe_topics();
1077        assert!(pending.contains(&"invalid.topic.BTCUSDT".to_string()));
1078    }
1079
1080    #[rstest]
1081    fn test_already_subscribed_error_confirms_topic() {
1082        let handler = create_test_handler();
1083        handler.subscriptions.mark_subscribe("tickers.ETHUSDT");
1084
1085        let resp = BybitWsResponse {
1086            op: Some(BybitWsOperation::Subscribe),
1087            topic: None,
1088            success: Some(false),
1089            conn_id: None,
1090            req_id: Some("tickers.ETHUSDT".to_string()),
1091            ret_code: Some(10001),
1092            ret_msg: Some("error:already subscribed,topic:tickers.ETHUSDT".to_string()),
1093        };
1094
1095        handler.handle_subscription_error(&resp);
1096
1097        let pending = handler.subscriptions.pending_subscribe_topics();
1098        assert!(!pending.contains(&"tickers.ETHUSDT".to_string()));
1099        let symbols = handler.subscriptions.confirmed();
1100        let entry = symbols
1101            .get(&Ustr::from("tickers"))
1102            .expect("channel present");
1103        assert!(entry.contains(&Ustr::from("ETHUSDT")));
1104    }
1105
1106    #[rstest]
1107    fn test_subscription_ack_without_req_id_confirms_all_pending() {
1108        let handler = create_test_handler();
1109        handler.subscriptions.mark_subscribe("orderbook.50.BTCUSDT");
1110        handler.subscriptions.mark_subscribe("publicTrade.BTCUSDT");
1111
1112        let ack = BybitWsSubscriptionMsg {
1113            success: true,
1114            op: BybitWsOperation::Subscribe,
1115            conn_id: None,
1116            req_id: None,
1117            ret_msg: None,
1118        };
1119
1120        handler.handle_subscription_ack(&ack);
1121
1122        // Both should be confirmed when no req_id
1123        assert!(handler.subscriptions.pending_subscribe_topics().is_empty());
1124    }
1125}
nautilus_bybit/websocket/handler.rs

nautilus_bybit/websocket/
handler.rs
