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

nautilus_bybit/websocket/
handler.rs
