nautilus_okx/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 OKX.
17//!
18//! The handler is a thin I/O boundary between the network layer and the client. It owns the
19//! `WebSocketClient`, deserializes raw venue messages into `OKXWsMessage` events, and handles
20//! subscription management, authentication, and retry logic.
21//!
22//! All domain parsing (venue types to Nautilus types) occurs outside the handler:
23//! - Data parsing in `PyOKXWebSocketClient` (uses an instruments cache)
24//! - Execution parsing in `execution.rs` (uses the system Cache)
25
26use std::{
27    collections::VecDeque,
28    fmt::Debug,
29    sync::{
30        Arc,
31        atomic::{AtomicBool, Ordering},
32    },
33};
34
35use nautilus_common::live::dst::time;
36use nautilus_core::{
37    AtomicTime,
38    string::secret::{REDACTED, SecretString},
39};
40pub use nautilus_live::book::snapshot::SnapshotGate;
41use nautilus_model::identifiers::ClientOrderId;
42use nautilus_network::{
43    RECONNECTED,
44    error::SendError,
45    retry::{RetryError, RetryManager, create_websocket_retry_manager},
46    websocket::{AuthTracker, SubscriptionState, TEXT_PING, TEXT_PONG, WebSocketClient},
47};
48use serde_json::{Map, Value};
49use tokio_tungstenite::tungstenite::Message;
50use tokio_util::sync::CancellationToken;
51use ustr::Ustr;
52
53use super::{
54    enums::{OKXSubscriptionEvent, OKXWsChannel, OKXWsOperation},
55    error::OKXWsError,
56    messages::{
57        OKXOrderMsg, OKXSubscription, OKXSubscriptionArg, OKXWebSocketArg, OKXWebSocketError,
58        OKXWsFrame, OKXWsMessage,
59    },
60    subscription::{topic_from_subscription_arg, topic_from_websocket_arg},
61};
62use crate::{
63    common::{
64        consts::{OKX_FIELD_SMSG, OKX_SUCCESS_CODE, should_retry_error_code},
65        enums::{OKXOrderStatus, OKXOrderType},
66        parse::prefer_rpi_response_fields,
67    },
68    websocket::client::OKX_RATE_LIMIT_KEY_SUBSCRIPTION,
69};
70
71/// Commands sent from the outer client to the inner message handler.
72pub enum HandlerCommand {
73    /// Set the `WebSocketClient` for the handler to use.
74    SetClient(WebSocketClient),
75    /// Disconnect the WebSocket connection.
76    Disconnect,
77    /// Send authentication payload to the WebSocket.
78    Authenticate { payload: SecretString },
79    /// Subscribe to the given channels.
80    Subscribe { args: Vec<OKXSubscriptionArg> },
81    /// Subscribes to a book and reports completion of the transport send.
82    SubscribeBook {
83        subscription: OKXSubscriptionArg,
84        cancel: CancellationToken,
85        gate: SnapshotGate,
86        completion: tokio::sync::oneshot::Sender<Result<(), OKXWsError>>,
87    },
88    /// Unsubscribe from the given channels.
89    Unsubscribe { args: Vec<OKXSubscriptionArg> },
90    /// Replaces a subscription without removing reconnect intent, unless canceled.
91    Resubscribe {
92        subscription: OKXSubscriptionArg,
93        cancel: CancellationToken,
94        gate: SnapshotGate,
95        completion: tokio::sync::oneshot::Sender<Result<(), OKXWsError>>,
96    },
97    /// Send a pre-serialized payload (used for order operations).
98    Send {
99        payload: String,
100        rate_limit_keys: Option<Vec<Ustr>>,
101        request_id: Option<String>,
102        client_order_ids: Vec<ClientOrderId>,
103        op: Option<OKXWsOperation>,
104    },
105}
106
107impl Debug for HandlerCommand {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        match self {
110            Self::SetClient(_) => f.write_str("SetClient"),
111            Self::Disconnect => f.write_str("Disconnect"),
112            Self::Authenticate { .. } => f
113                .debug_struct(stringify!(Authenticate))
114                .field("payload", &REDACTED)
115                .finish(),
116            Self::Subscribe { args } => f
117                .debug_struct(stringify!(Subscribe))
118                .field("args", args)
119                .finish(),
120            Self::SubscribeBook { subscription, .. } => {
121                f.debug_tuple("SubscribeBook").field(subscription).finish()
122            }
123            Self::Unsubscribe { args } => f
124                .debug_struct(stringify!(Unsubscribe))
125                .field("args", args)
126                .finish(),
127            Self::Resubscribe { subscription, .. } => {
128                f.debug_tuple("Resubscribe").field(subscription).finish()
129            }
130            Self::Send {
131                rate_limit_keys,
132                request_id,
133                client_order_ids,
134                op,
135                ..
136            } => f
137                .debug_struct(stringify!(Send))
138                .field("payload", &REDACTED)
139                .field("rate_limit_keys", rate_limit_keys)
140                .field("request_id", request_id)
141                .field("client_order_ids", client_order_ids)
142                .field("op", op)
143                .finish(),
144        }
145    }
146}
147
148pub(super) struct OKXWsFeedHandler {
149    clock: &'static AtomicTime,
150    signal: Arc<AtomicBool>,
151    inner: Option<WebSocketClient>,
152    cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
153    raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
154    out_tx: tokio::sync::mpsc::UnboundedSender<OKXWsMessage>,
155    auth_tracker: AuthTracker,
156    subscriptions_state: SubscriptionState,
157    retry_manager: RetryManager<OKXWsError>,
158    pending_messages: VecDeque<OKXWsMessage>,
159}
160
161impl OKXWsFeedHandler {
162    /// Creates a new [`OKXWsFeedHandler`] instance.
163    pub(super) fn new(
164        signal: Arc<AtomicBool>,
165        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
166        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
167        out_tx: tokio::sync::mpsc::UnboundedSender<OKXWsMessage>,
168        auth_tracker: AuthTracker,
169        subscriptions_state: SubscriptionState,
170        clock: &'static AtomicTime,
171    ) -> Self {
172        Self {
173            clock,
174            signal,
175            inner: None,
176            cmd_rx,
177            raw_rx,
178            out_tx,
179            auth_tracker,
180            subscriptions_state,
181            retry_manager: create_websocket_retry_manager(),
182            pending_messages: VecDeque::new(),
183        }
184    }
185
186    pub(super) fn is_stopped(&self) -> bool {
187        self.signal.load(Ordering::Acquire)
188    }
189
190    pub(super) fn send(&self, msg: OKXWsMessage) -> Result<(), ()> {
191        self.out_tx.send(msg).map_err(|_| ())
192    }
193
194    async fn send_with_retry(
195        &self,
196        payload: String,
197        rate_limit_keys: Option<&[Ustr]>,
198    ) -> Result<(), OKXWsError> {
199        self.send_secret_with_retry(payload.into(), rate_limit_keys)
200            .await
201    }
202
203    async fn send_secret_with_retry(
204        &self,
205        payload: SecretString,
206        rate_limit_keys: Option<&[Ustr]>,
207    ) -> Result<(), OKXWsError> {
208        if let Some(client) = &self.inner {
209            let keys_owned: Option<Vec<Ustr>> = rate_limit_keys.map(<[Ustr]>::to_vec);
210            self.retry_manager
211                .invocation(
212                    "websocket_send",
213                    || {
214                        let payload = payload.clone();
215                        let keys = keys_owned.clone();
216                        async move {
217                            client
218                                .send_text(payload.expose_secret().to_owned(), keys.as_deref())
219                                .await
220                                .map_err(OKXWsError::TransportSend)
221                        }
222                    },
223                    should_retry_replay_safe_error,
224                    create_okx_retry_error,
225                )
226                .execute()
227                .await
228        } else {
229            Err(OKXWsError::NoActiveClient)
230        }
231    }
232
233    async fn send_on_connection(
234        &self,
235        payload: String,
236        rate_limit_keys: Option<&[Ustr]>,
237    ) -> Result<(), OKXWsError> {
238        let client = self.inner.as_ref().ok_or(OKXWsError::NoActiveClient)?;
239        let connection_epoch = client.connection_epoch();
240        client
241            .send_text_on_connection(payload, rate_limit_keys, connection_epoch)
242            .await
243            .map_err(OKXWsError::TransportSend)
244    }
245
246    pub(super) async fn send_pong(&self) -> anyhow::Result<()> {
247        match self.send_on_connection(TEXT_PONG.to_string(), None).await {
248            Ok(()) => {
249                log::trace!("Sent pong response to OKX text ping");
250                Ok(())
251            }
252            Err(e) => Err(anyhow::anyhow!("Failed to send pong: {e}")),
253        }
254    }
255
256    pub(super) async fn next(&mut self) -> Option<OKXWsMessage> {
257        if let Some(message) = self.pending_messages.pop_front() {
258            return Some(message);
259        }
260
261        let mut poll_raw_next = false;
262
263        loop {
264            if self.signal.load(Ordering::Acquire) {
265                log::debug!("Stop signal received");
266                return None;
267            }
268
269            tokio::select! {
270                biased;
271                Some(cmd) = self.cmd_rx.recv(), if !poll_raw_next => {
272                    match cmd {
273                        HandlerCommand::SetClient(client) => {
274                            log::debug!("Handler received WebSocket client");
275                            self.inner = Some(client);
276                        }
277                        HandlerCommand::Disconnect => {
278                            log::debug!("Handler disconnecting WebSocket client");
279                            self.inner = None;
280                            return None;
281                        }
282                        HandlerCommand::Authenticate { payload } => {
283                            if let Err(e) = self.send_secret_with_retry(
284                                payload,
285                                Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()),
286                            ).await {
287                                log::error!(
288                                    "Failed to send authentication message after retries: error={e}"
289                                );
290                            }
291                        }
292                        HandlerCommand::Subscribe { args } => {
293                            if let Err(e) = self.handle_subscribe(args).await {
294                                log::error!("Failed to handle subscribe command: error={e}");
295                            }
296                        }
297                        HandlerCommand::SubscribeBook { subscription, cancel, gate, completion } => {
298                            let result = tokio::select! {
299                                biased;
300                                () = cancel.cancelled() => continue,
301                                result = async {
302                                    let client = self.inner.as_ref().ok_or(OKXWsError::NoActiveClient)?;
303                                    self.send_book_subscription(
304                                        OKXWsOperation::Subscribe,
305                                        subscription,
306                                        client.connection_epoch(),
307                                    ).await
308                                } => result,
309                            };
310
311                            if result.is_ok() {
312                                gate.open();
313                            }
314                            let _ = completion.send(result);
315                        }
316                        HandlerCommand::Unsubscribe { args } => {
317                            if let Err(e) = self.handle_unsubscribe(args).await {
318                                log::error!("Failed to handle unsubscribe command: error={e}");
319                            }
320                        }
321                        HandlerCommand::Resubscribe { subscription, cancel, gate, completion } => {
322                            let result = tokio::select! {
323                                biased;
324                                () = cancel.cancelled() => continue,
325                                result = async {
326                                    self.subscriptions_state.mark_failure(&topic_from_subscription_arg(&subscription));
327                                    let client = self.inner.as_ref().ok_or(OKXWsError::NoActiveClient)?;
328                                    let epoch = client.connection_epoch();
329                                    self.send_book_subscription(
330                                        OKXWsOperation::Unsubscribe,
331                                        subscription.clone(),
332                                        epoch,
333                                    ).await?;
334                                    self.send_book_subscription(
335                                        OKXWsOperation::Subscribe,
336                                        subscription,
337                                        epoch,
338                                    ).await
339                                } => result,
340                            };
341
342                            if result.is_ok() {
343                                gate.open();
344                            }
345                            let _ = completion.send(result);
346                        }
347                        HandlerCommand::Send {
348                            payload,
349                            rate_limit_keys,
350                            request_id,
351                            client_order_ids,
352                            op,
353                        } => {
354                            if let Err(e) = self.send_on_connection(
355                                payload,
356                                rate_limit_keys.as_deref(),
357                            ).await {
358                                log::debug!("Failed to send message: error={e}");
359
360                                if let Some(request_id) = request_id {
361                                    self.pending_messages.push_back(OKXWsMessage::SendFailed {
362                                        request_id,
363                                        client_order_ids,
364                                        op,
365                                        error: e,
366                                    });
367                                }
368                            }
369                        }
370                    }
371
372                    poll_raw_next = true;
373                }
374
375                msg = self.raw_rx.recv() => {
376                    let event = match msg {
377                        Some(msg) => match Self::parse_raw_message(msg) {
378                            Some(event) => event,
379                            None => continue,
380                        },
381                        None => {
382                            log::debug!("WebSocket stream closed");
383                            return None;
384                        }
385                    };
386
387                    match event {
388                        OKXWsFrame::Ping => {
389                            if let Err(e) = self.send_pong().await {
390                                log::warn!("Failed to send pong response: error={e}");
391                            }
392                        }
393                        OKXWsFrame::Login {
394                            code, msg, conn_id, ..
395                        } => {
396                            if code == OKX_SUCCESS_CODE {
397                                self.auth_tracker.succeed();
398                                return Some(OKXWsMessage::Authenticated);
399                            }
400
401                            self.auth_tracker.fail(msg.clone());
402
403                            let error = OKXWebSocketError {
404                                code,
405                                message: msg,
406                                conn_id: Some(conn_id),
407                                timestamp: self.clock.get_time_ns().as_u64(),
408                            };
409                            self.pending_messages.push_back(OKXWsMessage::Error(error));
410                        }
411                        OKXWsFrame::BookData { arg, action, data } => {
412                            return Some(OKXWsMessage::BookData { arg, action, data });
413                        }
414                        OKXWsFrame::RpiBookData { arg, action, data } => {
415                            return Some(OKXWsMessage::RpiBookData { arg, action, data });
416                        }
417                        OKXWsFrame::OrderResponse {
418                            id, op, code, msg, data,
419                        } => {
420                            return Some(OKXWsMessage::OrderResponse {
421                                id, op, code, msg, data,
422                            });
423                        }
424                        OKXWsFrame::Data { arg, data } => {
425                            if let Some(output) = self.route_data_message(arg, data) {
426                                return Some(output);
427                            }
428                        }
429                        OKXWsFrame::Error { arg, code, msg } => {
430                            let arg = arg.or_else(|| subscription_arg_from_error_message(&msg));
431                            if let Some(arg) = arg
432                                && self.handle_subscription_error(&arg, &code, &msg)
433                            {
434                                return Some(OKXWsMessage::SubscriptionFailed {
435                                    channel: arg.channel,
436                                    inst_id: arg.inst_id,
437                                    code,
438                                    msg,
439                                });
440                            }
441
442                            let error = OKXWebSocketError {
443                                code,
444                                message: msg,
445                                conn_id: None,
446                                timestamp: self.clock.get_time_ns().as_u64(),
447                            };
448                            return Some(OKXWsMessage::Error(error));
449                        }
450                        OKXWsFrame::Reconnected => {
451                            self.auth_tracker.invalidate();
452                            return Some(OKXWsMessage::Reconnected);
453                        }
454                        OKXWsFrame::Subscription {
455                            event, arg, code, msg,
456                            ..
457                        } => {
458                            let rejected = self
459                                .handle_subscription_ack(&event, &arg, code.as_deref(), msg.as_deref());
460
461                            if rejected {
462                                return Some(OKXWsMessage::SubscriptionFailed {
463                                    channel: arg.channel,
464                                    inst_id: arg.inst_id,
465                                    code: code.unwrap_or_default(),
466                                    msg: msg.unwrap_or_default(),
467                                });
468                            }
469                        }
470                        OKXWsFrame::ChannelConnCount { .. } => {}
471                    }
472                }
473
474                () = std::future::ready(()), if poll_raw_next => {
475                    poll_raw_next = false;
476                }
477
478                // Polled last so a ready message never registers and drops a timer
479                () = time::sleep(time::Duration::from_millis(100)) => {
480                    // Wake the loop to poll the stop signal while both channels are idle
481                }
482
483                else => {
484                    log::debug!("Handler shutting down: stream ended or command channel closed");
485                    return None;
486                }
487            }
488        }
489    }
490
491    fn route_data_message(&self, arg: OKXWebSocketArg, mut data: Value) -> Option<OKXWsMessage> {
492        let OKXWebSocketArg {
493            channel, inst_id, ..
494        } = arg;
495
496        match channel {
497            OKXWsChannel::Account => Some(OKXWsMessage::Account(data)),
498            OKXWsChannel::Positions => Some(OKXWsMessage::Positions(data)),
499            OKXWsChannel::Orders => {
500                parse_array_items(data, "orders", false).map(OKXWsMessage::Orders)
501            }
502            OKXWsChannel::SprdOrders => {
503                parse_array_items(data, "spread orders", false).map(OKXWsMessage::SpreadOrders)
504            }
505            OKXWsChannel::OrdersAlgo | OKXWsChannel::AlgoAdvance => {
506                parse_array_items(data, "algo orders", false).map(OKXWsMessage::AlgoOrders)
507            }
508            OKXWsChannel::LiquidationWarning => {
509                parse_array_items(data, "liquidation warnings", false)
510                    .map(OKXWsMessage::LiquidationWarnings)
511            }
512            OKXWsChannel::Instruments => {
513                prefer_rpi_response_fields(&mut data);
514                parse_array_items(data, "instruments", true).map(OKXWsMessage::Instruments)
515            }
516            _ => Some(OKXWsMessage::ChannelData {
517                channel,
518                inst_id,
519                data,
520            }),
521        }
522    }
523
524    fn handle_subscription_ack(
525        &self,
526        event: &OKXSubscriptionEvent,
527        arg: &OKXWebSocketArg,
528        code: Option<&str>,
529        msg: Option<&str>,
530    ) -> bool {
531        let topic = topic_from_websocket_arg(arg);
532        let success = code.is_none_or(|c| c == OKX_SUCCESS_CODE);
533
534        match event {
535            OKXSubscriptionEvent::Subscribe => {
536                if success {
537                    self.subscriptions_state.confirm_subscribe(&topic);
538                    false
539                } else {
540                    log::warn!(
541                        "Subscription failed: topic={topic:?}, error={msg:?}, code={code:?}"
542                    );
543                    self.subscriptions_state.mark_failure(&topic);
544                    true
545                }
546            }
547            OKXSubscriptionEvent::Unsubscribe => {
548                if success {
549                    self.subscriptions_state.confirm_unsubscribe(&topic);
550                } else {
551                    log::warn!(
552                        "Unsubscription failed - restoring subscription: \
553                         topic={topic:?}, error={msg:?}, code={code:?}"
554                    );
555                    self.subscriptions_state.confirm_unsubscribe(&topic);
556                    self.subscriptions_state.mark_subscribe(&topic);
557                    self.subscriptions_state.confirm_subscribe(&topic);
558                }
559                false
560            }
561        }
562    }
563
564    fn handle_subscription_error(&self, arg: &OKXWebSocketArg, code: &str, msg: &str) -> bool {
565        let topic = topic_from_websocket_arg(arg);
566        let event = if self
567            .subscriptions_state
568            .pending_unsubscribe_topics()
569            .iter()
570            .any(|pending| pending == &topic)
571        {
572            OKXSubscriptionEvent::Unsubscribe
573        } else if self
574            .subscriptions_state
575            .pending_subscribe_topics()
576            .iter()
577            .any(|pending| pending == &topic)
578            || (arg.channel.is_book()
579                && self
580                    .subscriptions_state
581                    .all_topics()
582                    .iter()
583                    .any(|tracked| tracked == &topic))
584        {
585            OKXSubscriptionEvent::Subscribe
586        } else {
587            return false;
588        };
589
590        self.handle_subscription_ack(&event, arg, Some(code), Some(msg))
591    }
592
593    async fn send_book_subscription(
594        &self,
595        op: OKXWsOperation,
596        subscription: OKXSubscriptionArg,
597        connection_epoch: u64,
598    ) -> Result<(), OKXWsError> {
599        let client = self.inner.as_ref().ok_or(OKXWsError::NoActiveClient)?;
600        let message = OKXSubscription {
601            op,
602            args: vec![subscription],
603        };
604        let payload =
605            serde_json::to_string(&message).map_err(|e| OKXWsError::ClientError(e.to_string()))?;
606
607        client
608            .send_text_on_connection(
609                payload,
610                Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()),
611                connection_epoch,
612            )
613            .await
614            .map_err(OKXWsError::TransportSend)
615    }
616
617    async fn handle_subscribe(&self, args: Vec<OKXSubscriptionArg>) -> anyhow::Result<()> {
618        for arg in &args {
619            log::debug!(
620                "Subscribing to channel: channel={:?}, inst_id={:?}",
621                arg.channel,
622                arg.inst_id
623            );
624        }
625
626        let message = OKXSubscription {
627            op: OKXWsOperation::Subscribe,
628            args,
629        };
630
631        let json_txt = serde_json::to_string(&message)
632            .map_err(|e| anyhow::anyhow!("Failed to serialize subscription: {e}"))?;
633
634        self.send_with_retry(json_txt, Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
635            .await
636            .map_err(|e| anyhow::anyhow!("Failed to send subscription after retries: {e}"))?;
637        Ok(())
638    }
639
640    async fn handle_unsubscribe(&self, args: Vec<OKXSubscriptionArg>) -> anyhow::Result<()> {
641        for arg in &args {
642            log::debug!(
643                "Unsubscribing from channel: channel={:?}, inst_id={:?}",
644                arg.channel,
645                arg.inst_id
646            );
647        }
648
649        let message = OKXSubscription {
650            op: OKXWsOperation::Unsubscribe,
651            args,
652        };
653
654        let json_txt = serde_json::to_string(&message)
655            .map_err(|e| anyhow::anyhow!("Failed to serialize unsubscription: {e}"))?;
656
657        self.send_with_retry(json_txt, Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
658            .await
659            .map_err(|e| anyhow::anyhow!("Failed to send unsubscription after retries: {e}"))?;
660        Ok(())
661    }
662
663    pub(crate) fn parse_raw_message(
664        msg: tokio_tungstenite::tungstenite::Message,
665    ) -> Option<OKXWsFrame> {
666        match msg {
667            tokio_tungstenite::tungstenite::Message::Text(text) => {
668                if text == TEXT_PONG {
669                    log::trace!("Received pong from OKX");
670                    return None;
671                }
672
673                if text == TEXT_PING {
674                    log::trace!("Received ping from OKX (text)");
675                    return Some(OKXWsFrame::Ping);
676                }
677
678                if text == RECONNECTED {
679                    log::debug!("Received WebSocket reconnection signal");
680                    return Some(OKXWsFrame::Reconnected);
681                }
682                log::trace!("Received WebSocket message: {text}");
683
684                match serde_json::from_str(&text) {
685                    Ok(ws_event) => match &ws_event {
686                        OKXWsFrame::Error { arg, code, msg } => {
687                            let level = error_frame_level(arg.as_ref(), code, msg);
688                            log::log!(level, "WebSocket error: {code} - {msg}");
689                            Some(ws_event)
690                        }
691                        OKXWsFrame::Login {
692                            event,
693                            code,
694                            msg,
695                            conn_id,
696                        } => {
697                            if code == OKX_SUCCESS_CODE {
698                                log::debug!("WebSocket authenticated: conn_id={conn_id}");
699                            } else {
700                                log::error!(
701                                    "WebSocket authentication failed: \
702                                     event={event}, code={code}, error={msg}"
703                                );
704                            }
705                            Some(ws_event)
706                        }
707                        OKXWsFrame::Subscription {
708                            event,
709                            arg,
710                            conn_id,
711                            ..
712                        } => {
713                            let channel_str = serde_json::to_string(&arg.channel)
714                                .expect("Invalid OKX websocket channel")
715                                .trim_matches('"')
716                                .to_string();
717                            log::debug!("{event}d: channel={channel_str}, conn_id={conn_id}");
718                            Some(ws_event)
719                        }
720                        OKXWsFrame::ChannelConnCount {
721                            channel,
722                            conn_count,
723                            conn_id,
724                            ..
725                        } => {
726                            let channel_str = serde_json::to_string(channel)
727                                .expect("Invalid OKX websocket channel")
728                                .trim_matches('"')
729                                .to_string();
730                            log::debug!(
731                                "Channel connection status: \
732                                 channel={channel_str}, connections={conn_count}, conn_id={conn_id}",
733                            );
734                            None
735                        }
736                        OKXWsFrame::Ping => {
737                            log::trace!("Ignoring ping event parsed from text payload");
738                            None
739                        }
740                        OKXWsFrame::Data { .. }
741                        | OKXWsFrame::BookData { .. }
742                        | OKXWsFrame::RpiBookData { .. } => Some(ws_event),
743                        OKXWsFrame::OrderResponse {
744                            id, op, code, data, ..
745                        } => {
746                            if code == OKX_SUCCESS_CODE {
747                                log::debug!(
748                                    "Order operation successful: id={id:?}, op={op}, code={code}"
749                                );
750
751                                if let Some(order_data) = data.first() {
752                                    let success_msg = order_data
753                                        .get(OKX_FIELD_SMSG)
754                                        .and_then(|s| s.as_str())
755                                        .unwrap_or("Order operation successful");
756                                    log::debug!("Order success details: {success_msg}");
757                                }
758                            }
759                            Some(ws_event)
760                        }
761                        OKXWsFrame::Reconnected => {
762                            log::warn!("Unexpected Reconnected event from deserialization");
763                            None
764                        }
765                    },
766                    Err(e) => {
767                        log::error!("Failed to parse message: {e}: {text}");
768                        None
769                    }
770                }
771            }
772            Message::Ping(_payload) => {
773                log::trace!("Received binary ping frame from OKX");
774                Some(OKXWsFrame::Ping)
775            }
776            Message::Pong(payload) => {
777                log::trace!("Received pong frame from OKX ({} bytes)", payload.len());
778                None
779            }
780            Message::Binary(msg) => {
781                log::debug!("Raw binary frame ({} bytes)", msg.len());
782                log::trace!("Raw binary: {msg:?}");
783                None
784            }
785            Message::Close(_) => {
786                log::debug!("Received close message");
787                None
788            }
789            msg => {
790                log::warn!("Unexpected message: {msg}");
791                None
792            }
793        }
794    }
795}
796
797// Book recovery retries a rejected book subscription and reports its own failure
798fn error_frame_level(arg: Option<&OKXWebSocketArg>, code: &str, msg: &str) -> log::Level {
799    let book = match arg {
800        Some(arg) => arg.channel.is_book(),
801        None => subscription_arg_from_error_message(msg).is_some_and(|arg| arg.channel.is_book()),
802    };
803
804    if should_retry_error_code(code) || book {
805        log::Level::Warn
806    } else {
807        log::Level::Error
808    }
809}
810
811fn subscription_arg_from_error_message(msg: &str) -> Option<OKXWebSocketArg> {
812    let descriptor = msg
813        .strip_prefix("Wrong URL or channel:")?
814        .split_whitespace()
815        .next()?;
816    let mut fields = descriptor.split(',');
817    let channel = fields.next()?;
818    let mut arg = Map::new();
819    arg.insert("channel".to_string(), Value::String(channel.to_string()));
820
821    for field in fields {
822        let (key, value) = field.split_once(':')?;
823        if !matches!(key, "instId" | "sprdId" | "instType" | "instFamily") {
824            return None;
825        }
826        arg.insert(key.to_string(), Value::String(value.to_string()));
827    }
828
829    serde_json::from_value(Value::Object(arg)).ok()
830}
831
832/// Returns `true` when an OKX WebSocket order message represents a post-only auto-cancel.
833pub fn is_post_only_auto_cancel(msg: &OKXOrderMsg) -> bool {
834    use crate::common::{consts::OKX_POST_ONLY_CANCEL_SOURCE, enums::OKXOrderStatus};
835
836    if msg.state != OKXOrderStatus::Canceled {
837        return false;
838    }
839
840    let cancel_source_matches = matches!(
841        msg.cancel_source.as_deref(),
842        Some(source) if source == OKX_POST_ONLY_CANCEL_SOURCE
843    );
844
845    let reason_matches = matches!(
846        msg.cancel_source_reason.as_deref(),
847        Some(reason) if reason.contains("POST_ONLY")
848    );
849
850    if !(cancel_source_matches || reason_matches) {
851        return false;
852    }
853
854    msg.acc_fill_sz
855        .as_ref()
856        .is_none_or(|filled| filled == "0" || filled.is_empty())
857}
858
859/// Returns `true` when an RPI order update is canceled without any fill.
860pub fn is_unfilled_rpi_cancel(msg: &OKXOrderMsg) -> bool {
861    msg.ord_type == OKXOrderType::Rpi
862        && msg.state == OKXOrderStatus::Canceled
863        && msg
864            .acc_fill_sz
865            .as_ref()
866            .is_none_or(|filled| filled == "0" || filled.is_empty())
867}
868
869// Per-item deserialization so one malformed entry does not drop the batch.
870fn parse_array_items<T: serde::de::DeserializeOwned>(
871    data: Value,
872    label: &str,
873    warn_on_parse_error: bool,
874) -> Option<Vec<T>> {
875    let Value::Array(items) = data else {
876        if warn_on_parse_error {
877            log::warn!("Expected {label} payload to be a JSON array");
878        } else {
879            log::error!("Expected {label} payload to be a JSON array");
880        }
881        return None;
882    };
883
884    let mut parsed = Vec::with_capacity(items.len());
885    for (idx, item) in items.into_iter().enumerate() {
886        match serde_json::from_value::<T>(item) {
887            Ok(value) => parsed.push(value),
888            Err(e) => {
889                if warn_on_parse_error {
890                    log::warn!("Failed to parse {label} item at index {idx}: {e}");
891                } else {
892                    log::error!("Failed to parse {label} item at index {idx}: {e}");
893                }
894            }
895        }
896    }
897
898    if parsed.is_empty() {
899        None
900    } else {
901        Some(parsed)
902    }
903}
904
905fn should_retry_replay_safe_error(error: &OKXWsError) -> bool {
906    match error {
907        OKXWsError::OkxError { error_code, .. } => should_retry_error_code(error_code),
908        OKXWsError::TransportSend(SendError::Timeout | SendError::ConnectionChanged)
909        | OKXWsError::TungsteniteError(_)
910        | OKXWsError::OperationTimeout { .. } => true,
911        OKXWsError::AuthenticationError(_)
912        | OKXWsError::JsonError(_)
913        | OKXWsError::ParsingError(_)
914        | OKXWsError::ClientError(_)
915        | OKXWsError::NoActiveClient
916        | OKXWsError::HandlerUnavailable(_)
917        | OKXWsError::TransportSend(
918            SendError::InvalidInput(_)
919            | SendError::BufferFull
920            | SendError::Closed
921            | SendError::WriteTimeout
922            | SendError::BrokenPipe(_),
923        )
924        | OKXWsError::SendFailed(_) => false,
925    }
926}
927
928fn create_okx_retry_error(error: RetryError) -> OKXWsError {
929    match error {
930        RetryError::OperationTimeout { timeout_ms } => OKXWsError::OperationTimeout { timeout_ms },
931        RetryError::InvalidConfiguration { message } => OKXWsError::ClientError(message),
932        RetryError::Canceled => {
933            OKXWsError::SendFailed("Adapter disconnecting or shutting down".to_string())
934        }
935        error @ RetryError::ElapsedBudgetExceeded { .. } => {
936            OKXWsError::SendFailed(error.to_string())
937        }
938    }
939}
940
941#[cfg(test)]
942mod tests {
943    use std::{
944        num::NonZeroU32,
945        sync::{Arc, atomic::AtomicBool},
946        time::Duration,
947    };
948
949    use futures_util::{SinkExt, StreamExt};
950    use nautilus_core::{collections::AtomicMap, time::get_atomic_clock_realtime};
951    use nautilus_model::identifiers::InstrumentId;
952    use nautilus_network::{
953        ratelimiter::{RateLimiter, quota::Quota},
954        websocket::{AuthTracker, SubscriptionState, WebSocketConfig, channel_message_handler},
955    };
956    use rstest::rstest;
957    use serde_json::json;
958
959    use super::*;
960    use crate::{
961        book::{BookChannelScope, BookSequenceOutcome, sync::BookSyncTracker},
962        common::{
963            consts::OKX_WS_TOPIC_DELIMITER,
964            enums::{OKXBookAction, OKXBookChannel, OKXRpiPermission},
965            testing::load_test_json,
966        },
967    };
968
969    fn create_handler() -> OKXWsFeedHandler {
970        let signal = Arc::new(AtomicBool::new(false));
971        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
972        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
973        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
974
975        OKXWsFeedHandler::new(
976            signal,
977            cmd_rx,
978            raw_rx,
979            out_tx,
980            AuthTracker::new(),
981            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
982            get_atomic_clock_realtime(),
983        )
984    }
985
986    #[rstest]
987    fn test_command_debug_redacts_payloads() {
988        let payload = "authentication-secret";
989        let authenticate = HandlerCommand::Authenticate {
990            payload: SecretString::from(payload.to_string()),
991        };
992        let send = HandlerCommand::Send {
993            payload: payload.to_string(),
994            rate_limit_keys: None,
995            request_id: None,
996            client_order_ids: Vec::new(),
997            op: None,
998        };
999
1000        let debug = format!("{authenticate:?} {send:?}");
1001
1002        assert!(debug.contains(REDACTED));
1003        assert!(!debug.contains(payload));
1004    }
1005
1006    #[tokio::test]
1007    async fn test_next_polls_raw_after_one_ready_command() {
1008        let signal = Arc::new(AtomicBool::new(false));
1009        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1010        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1011        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1012        let mut handler = OKXWsFeedHandler::new(
1013            signal,
1014            cmd_rx,
1015            raw_rx,
1016            out_tx,
1017            AuthTracker::new(),
1018            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1019            get_atomic_clock_realtime(),
1020        );
1021
1022        for _ in 0..3 {
1023            cmd_tx
1024                .send(HandlerCommand::Subscribe { args: Vec::new() })
1025                .unwrap();
1026        }
1027        raw_tx
1028            .send(Message::Text(RECONNECTED.to_string().into()))
1029            .unwrap();
1030
1031        let message = handler.next().await;
1032
1033        assert!(matches!(message, Some(OKXWsMessage::Reconnected)));
1034        assert_eq!(handler.cmd_rx.len(), 2);
1035    }
1036
1037    #[rstest]
1038    #[case::sent(false)]
1039    #[case::canceled(true)]
1040    #[tokio::test]
1041    async fn initial_book_completion_waits_for_send(#[case] canceled: bool) {
1042        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1043        let url = format!("ws://{}/", listener.local_addr().unwrap());
1044        let (wire_tx, mut wire_rx) = tokio::sync::mpsc::unbounded_channel();
1045
1046        let server = tokio::spawn(async move {
1047            let (socket, _) = listener.accept().await.unwrap();
1048            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1049            while let Some(Ok(frame)) = socket.next().await {
1050                wire_tx.send(frame).unwrap();
1051            }
1052        });
1053
1054        let (message_handler, raw_rx) = channel_message_handler();
1055        let client = WebSocketClient::builder()
1056            .config(WebSocketConfig::builder().url(url).build().unwrap())
1057            .message_handler(message_handler)
1058            .default_quota(Quota::with_period(Duration::from_secs(10)).unwrap())
1059            .connect()
1060            .await
1061            .unwrap();
1062        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1063        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1064
1065        let mut handler = OKXWsFeedHandler::new(
1066            Arc::new(AtomicBool::new(false)),
1067            cmd_rx,
1068            raw_rx,
1069            out_tx,
1070            AuthTracker::new(),
1071            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1072            get_atomic_clock_realtime(),
1073        );
1074        handler.inner = Some(client);
1075        cmd_tx
1076            .send(HandlerCommand::Subscribe { args: Vec::new() })
1077            .unwrap();
1078        let cancel = CancellationToken::new();
1079        let gate = SnapshotGate::default();
1080        gate.lock().close();
1081        let (completion, mut result) = tokio::sync::oneshot::channel();
1082        cmd_tx
1083            .send(HandlerCommand::SubscribeBook {
1084                subscription: OKXSubscriptionArg {
1085                    channel: OKXWsChannel::Books,
1086                    inst_id: Some(Ustr::from("BTC-USDT")),
1087                    inst_type: None,
1088                    inst_family: None,
1089                },
1090                cancel: cancel.clone(),
1091                gate: gate.clone(),
1092                completion,
1093            })
1094            .unwrap();
1095
1096        let running = tokio::spawn(async move { handler.next().await });
1097        let first = tokio::time::timeout(Duration::from_secs(3), wire_rx.recv())
1098            .await
1099            .unwrap()
1100            .unwrap();
1101        let first: Value = serde_json::from_str(first.to_text().unwrap()).unwrap();
1102        assert_eq!(first, serde_json::json!({"op": "subscribe", "args": []}));
1103        assert!(matches!(
1104            result.try_recv(),
1105            Err(tokio::sync::oneshot::error::TryRecvError::Empty)
1106        ));
1107
1108        if canceled {
1109            cancel.cancel();
1110        }
1111
1112        tokio::time::pause();
1113        tokio::time::advance(Duration::from_secs(10)).await;
1114        tokio::time::resume();
1115        let result = tokio::time::timeout(Duration::from_secs(3), result)
1116            .await
1117            .unwrap();
1118
1119        if canceled {
1120            assert!(gate.lock().is_closed());
1121            assert!(result.is_err());
1122            assert!(
1123                tokio::time::timeout(Duration::from_millis(100), wire_rx.recv())
1124                    .await
1125                    .is_err()
1126            );
1127        } else {
1128            result.unwrap().unwrap();
1129            assert!(!gate.lock().is_closed());
1130            let frame = tokio::time::timeout(Duration::from_secs(3), wire_rx.recv())
1131                .await
1132                .unwrap()
1133                .unwrap();
1134            let frame: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
1135            assert_eq!(
1136                frame,
1137                serde_json::json!({"op": "subscribe", "args": [{"channel": "books", "instId": "BTC-USDT"}]})
1138            );
1139        }
1140
1141        running.abort();
1142        server.abort();
1143    }
1144
1145    #[tokio::test]
1146    async fn book_subscription_rejects_changed_connection() {
1147        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1148        let url = format!("ws://{}/", listener.local_addr().unwrap());
1149        let (wire_tx, mut wire_rx) = tokio::sync::mpsc::unbounded_channel();
1150
1151        let server = tokio::spawn(async move {
1152            let (socket, _) = listener.accept().await.unwrap();
1153            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1154            while let Some(Ok(frame)) = socket.next().await {
1155                wire_tx.send(frame).unwrap();
1156            }
1157        });
1158        let (message_handler, raw_rx) = channel_message_handler();
1159        let client = WebSocketClient::builder()
1160            .config(WebSocketConfig::builder().url(url).build().unwrap())
1161            .message_handler(message_handler)
1162            .connect()
1163            .await
1164            .unwrap();
1165        let wrong_epoch = client.connection_epoch() + 1;
1166        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1167        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1168        let mut handler = OKXWsFeedHandler::new(
1169            Arc::new(AtomicBool::new(false)),
1170            cmd_rx,
1171            raw_rx,
1172            out_tx,
1173            AuthTracker::new(),
1174            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1175            get_atomic_clock_realtime(),
1176        );
1177        handler.inner = Some(client);
1178        let result = handler
1179            .send_book_subscription(
1180                OKXWsOperation::Subscribe,
1181                OKXSubscriptionArg {
1182                    channel: OKXWsChannel::Books,
1183                    inst_id: Some(Ustr::from("BTC-USDT")),
1184                    inst_type: None,
1185                    inst_family: None,
1186                },
1187                wrong_epoch,
1188            )
1189            .await;
1190
1191        assert!(matches!(
1192            result,
1193            Err(OKXWsError::TransportSend(SendError::ConnectionChanged))
1194        ));
1195        assert!(matches!(
1196            wire_rx.try_recv(),
1197            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1198        ));
1199        handler.inner.as_ref().unwrap().disconnect().await;
1200        server.abort();
1201    }
1202
1203    #[tokio::test]
1204    async fn recovery_cancel_between_sends_keeps_snapshot_gate_closed() {
1205        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1206        let url = format!("ws://{}/", listener.local_addr().unwrap());
1207        let (received, first_frame) = tokio::sync::oneshot::channel();
1208
1209        let server = tokio::spawn(async move {
1210            let (socket, _) = listener.accept().await.unwrap();
1211            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1212            let frame = socket.next().await.unwrap().unwrap();
1213            received.send(frame).unwrap();
1214
1215            while socket.next().await.is_some() {}
1216        });
1217
1218        let (message_handler, raw_rx) = channel_message_handler();
1219        let client = WebSocketClient::builder()
1220            .config(WebSocketConfig::builder().url(url).build().unwrap())
1221            .message_handler(message_handler)
1222            .default_quota(Quota::per_hour(NonZeroU32::new(1).unwrap()))
1223            .connect()
1224            .await
1225            .unwrap();
1226        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1227        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1228
1229        let mut handler = OKXWsFeedHandler::new(
1230            Arc::new(AtomicBool::new(false)),
1231            cmd_rx,
1232            raw_rx,
1233            out_tx,
1234            AuthTracker::new(),
1235            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1236            get_atomic_clock_realtime(),
1237        );
1238        handler.inner = Some(client);
1239        let tracker = BookSyncTracker::default();
1240        let instrument_id = InstrumentId::from("BTC-USDT.OKX");
1241        tracker.record_subscription(instrument_id, time::Instant::now(), SnapshotGate::default());
1242        let recovery = tracker.claim_recovery(instrument_id).unwrap();
1243        assert!(recovery.begin_replacement());
1244        let cancel = recovery.cancellation.child_token();
1245        let (completion, result) = tokio::sync::oneshot::channel();
1246        cmd_tx
1247            .send(HandlerCommand::Resubscribe {
1248                subscription: OKXSubscriptionArg {
1249                    channel: OKXWsChannel::Books,
1250                    inst_id: Some(Ustr::from("BTC-USDT")),
1251                    inst_type: None,
1252                    inst_family: None,
1253                },
1254                cancel: cancel.clone(),
1255                gate: recovery.gate.clone(),
1256                completion,
1257            })
1258            .unwrap();
1259
1260        let running = tokio::spawn(async move { handler.next().await });
1261        let frame = tokio::time::timeout(Duration::from_secs(3), first_frame)
1262            .await
1263            .unwrap()
1264            .unwrap();
1265        let request: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
1266        assert_eq!(request["op"], "unsubscribe");
1267        assert_eq!(
1268            tracker.validate_sequence(instrument_id, true, &[(Some(-1), 42)], time::Instant::now()),
1269            BookSequenceOutcome::Suppress
1270        );
1271        cancel.cancel();
1272        assert!(
1273            tokio::time::timeout(Duration::from_secs(1), result)
1274                .await
1275                .unwrap()
1276                .is_err()
1277        );
1278
1279        assert!(recovery.gate.lock().is_closed());
1280        assert_eq!(
1281            tracker.validate_sequence(instrument_id, true, &[(Some(-1), 43)], time::Instant::now()),
1282            BookSequenceOutcome::Suppress
1283        );
1284        cmd_tx.send(HandlerCommand::Disconnect).unwrap();
1285        tokio::time::timeout(Duration::from_secs(3), running)
1286            .await
1287            .unwrap()
1288            .unwrap();
1289        server.abort();
1290    }
1291
1292    #[rstest]
1293    #[case::disabled_deadline(Duration::ZERO)]
1294    #[case::snapshot_deadline(Duration::from_secs(3))]
1295    #[tokio::test]
1296    async fn reconnect_replay_survives_inflight_recovery(#[case] snapshot_timeout: Duration) {
1297        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1298        let url = format!("ws://{}/", listener.local_addr().unwrap());
1299        let (wire_tx, mut wire_rx) = tokio::sync::mpsc::unbounded_channel();
1300
1301        let server = tokio::spawn(async move {
1302            let (socket, _) = listener.accept().await.unwrap();
1303            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1304
1305            while let Some(Ok(Message::Text(text))) = socket.next().await {
1306                let request: Value = serde_json::from_str(&text).unwrap();
1307                wire_tx
1308                    .send(request["op"].as_str().unwrap().to_owned())
1309                    .unwrap();
1310
1311                if request["op"] == "subscribe" {
1312                    for (fixture, sequence, previous) in [
1313                        ("ws_books_snapshot.json", 100, -1),
1314                        ("ws_books_update.json", 101, 100),
1315                    ] {
1316                        let mut frame: Value =
1317                            serde_json::from_str(&load_test_json(fixture)).unwrap();
1318                        frame["data"][0]["seqId"] = json!(sequence);
1319                        frame["data"][0]["prevSeqId"] = json!(previous);
1320                        socket
1321                            .send(Message::Text(frame.to_string().into()))
1322                            .await
1323                            .unwrap();
1324                    }
1325                }
1326            }
1327        });
1328
1329        // Replay and unsubscribe consume the burst; the replacement subscribe cannot send
1330        let limiter = Arc::new(RateLimiter::new_with_quota(
1331            Some(
1332                Quota::with_period(Duration::from_secs(10))
1333                    .unwrap()
1334                    .allow_burst(NonZeroU32::new(2).unwrap()),
1335            ),
1336            Vec::new(),
1337        ));
1338        let (message_handler, raw_rx) = channel_message_handler();
1339        let client = WebSocketClient::builder()
1340            .config(WebSocketConfig::builder().url(url).build().unwrap())
1341            .message_handler(message_handler)
1342            .rate_limiter(Arc::clone(&limiter))
1343            .connect()
1344            .await
1345            .unwrap();
1346        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1347        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1348
1349        let mut handler = OKXWsFeedHandler::new(
1350            Arc::new(AtomicBool::new(false)),
1351            cmd_rx,
1352            raw_rx,
1353            out_tx,
1354            AuthTracker::new(),
1355            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1356            get_atomic_clock_realtime(),
1357        );
1358        handler.inner = Some(client);
1359        let (messages_tx, mut messages_rx) = tokio::sync::mpsc::unbounded_channel();
1360
1361        let running = tokio::spawn(async move {
1362            while let Some(message) = handler.next().await {
1363                messages_tx.send(message).unwrap();
1364            }
1365        });
1366
1367        let instrument_id = InstrumentId::from("BTC-USDT.OKX");
1368        let channels = AtomicMap::new();
1369        channels.insert(instrument_id, OKXBookChannel::Book);
1370        let tracker = BookSyncTracker::default();
1371        tracker.record_subscription(instrument_id, time::Instant::now(), SnapshotGate::default());
1372        let recovery = tracker.claim_recovery(instrument_id).unwrap();
1373        assert!(recovery.begin_replacement());
1374
1375        let arg = OKXSubscriptionArg {
1376            channel: OKXWsChannel::Books,
1377            inst_id: Some(Ustr::from("BTC-USDT")),
1378            inst_type: None,
1379            inst_family: None,
1380        };
1381
1382        cmd_tx
1383            .send(HandlerCommand::Subscribe {
1384                args: vec![arg.clone()],
1385            })
1386            .unwrap();
1387
1388        assert_eq!(
1389            tokio::time::timeout(Duration::from_secs(3), wire_rx.recv())
1390                .await
1391                .unwrap()
1392                .unwrap(),
1393            "subscribe"
1394        );
1395
1396        for action in [OKXBookAction::Snapshot, OKXBookAction::Update] {
1397            let message = tokio::time::timeout(Duration::from_secs(3), messages_rx.recv())
1398                .await
1399                .unwrap()
1400                .unwrap();
1401            assert!(
1402                matches!(message, OKXWsMessage::BookData { action: received, .. } if received == action)
1403            );
1404        }
1405
1406        let (completion, result) = tokio::sync::oneshot::channel();
1407        cmd_tx
1408            .send(HandlerCommand::Resubscribe {
1409                subscription: arg.clone(),
1410                cancel: recovery.cancellation.child_token(),
1411                gate: recovery.gate.clone(),
1412                completion,
1413            })
1414            .unwrap();
1415
1416        assert_eq!(
1417            tokio::time::timeout(Duration::from_secs(3), wire_rx.recv())
1418                .await
1419                .unwrap()
1420                .unwrap(),
1421            "unsubscribe"
1422        );
1423
1424        // Paused time keeps the subscribe blocked until the reset completes
1425        tokio::time::pause();
1426        tracker.reset_sequences(&channels, BookChannelScope::Public);
1427
1428        if !snapshot_timeout.is_zero() {
1429            tracker.seed_pending_snapshots(
1430                &channels,
1431                BookChannelScope::Public,
1432                snapshot_timeout,
1433                time::Instant::now(),
1434            );
1435        }
1436
1437        tokio::time::advance(Duration::from_secs(10)).await;
1438        tokio::time::resume();
1439        let _ = tokio::time::timeout(Duration::from_secs(3), result)
1440            .await
1441            .expect("recovery command completes after reconnect");
1442
1443        let resumed = tokio::time::timeout(Duration::from_secs(3), async {
1444            for (action, sequence) in [(OKXBookAction::Snapshot, 100), (OKXBookAction::Update, 101)]
1445            {
1446                let Some(OKXWsMessage::BookData {
1447                    action: received,
1448                    data,
1449                    ..
1450                }) = messages_rx.recv().await
1451                else {
1452                    panic!("expected book data after reconnect");
1453                };
1454
1455                assert_eq!(received, action);
1456                assert_eq!(data[0].seq_id, sequence);
1457                assert_eq!(
1458                    tracker.validate_sequence(
1459                        instrument_id,
1460                        received == OKXBookAction::Snapshot,
1461                        &[(data[0].prev_seq_id, data[0].seq_id)],
1462                        time::Instant::now(),
1463                    ),
1464                    BookSequenceOutcome::Accept,
1465                );
1466            }
1467        })
1468        .await;
1469
1470        // A successful manual subscribe isolates a failure to automatic replay
1471        if resumed.is_err() {
1472            cmd_tx
1473                .send(HandlerCommand::Subscribe { args: vec![arg] })
1474                .unwrap();
1475
1476            for action in [OKXBookAction::Snapshot, OKXBookAction::Update] {
1477                let message = tokio::time::timeout(Duration::from_secs(3), messages_rx.recv())
1478                    .await
1479                    .unwrap()
1480                    .unwrap();
1481                assert!(
1482                    matches!(message, OKXWsMessage::BookData { action: received, .. } if received == action)
1483                );
1484            }
1485        }
1486
1487        cmd_tx.send(HandlerCommand::Disconnect).unwrap();
1488        tokio::time::timeout(Duration::from_secs(3), running)
1489            .await
1490            .unwrap()
1491            .unwrap();
1492        server.abort();
1493
1494        assert!(
1495            resumed.is_ok(),
1496            "reconnect must preserve an in-flight recovery and resume book output"
1497        );
1498    }
1499
1500    #[rstest]
1501    fn test_should_retry_typed_transport_and_timeout_errors() {
1502        assert!(!should_retry_replay_safe_error(&OKXWsError::TransportSend(
1503            SendError::BufferFull
1504        )));
1505        assert!(should_retry_replay_safe_error(&OKXWsError::TransportSend(
1506            SendError::Timeout
1507        )));
1508        assert!(should_retry_replay_safe_error(&OKXWsError::TransportSend(
1509            SendError::ConnectionChanged
1510        )));
1511        assert!(!should_retry_replay_safe_error(&OKXWsError::TransportSend(
1512            SendError::WriteTimeout
1513        )));
1514        assert!(!should_retry_replay_safe_error(&OKXWsError::TransportSend(
1515            SendError::BrokenPipe("connection reset".to_string())
1516        )));
1517        assert!(should_retry_replay_safe_error(
1518            &OKXWsError::OperationTimeout { timeout_ms: 1_000 }
1519        ));
1520        assert!(!should_retry_replay_safe_error(&OKXWsError::NoActiveClient));
1521        assert!(!should_retry_replay_safe_error(
1522            &OKXWsError::HandlerUnavailable("closed".to_string())
1523        ));
1524    }
1525
1526    #[rstest]
1527    fn test_retryability_uses_websocket_error_type_not_message() {
1528        let message = "connection reset".to_string();
1529        let temporary = OKXWsError::OkxError {
1530            error_code: "50011".to_string(),
1531            message: message.clone(),
1532        };
1533        let permanent = OKXWsError::ClientError(message.clone());
1534        let ambiguous = OKXWsError::SendFailed(message);
1535
1536        assert!(should_retry_replay_safe_error(&temporary));
1537        assert!(!should_retry_replay_safe_error(&permanent));
1538        assert!(!should_retry_replay_safe_error(&ambiguous));
1539    }
1540
1541    #[rstest]
1542    fn test_subscription_error_restores_failed_unsubscribe() {
1543        let handler = create_handler();
1544        let arg = OKXWebSocketArg {
1545            channel: OKXWsChannel::Books,
1546            inst_id: Some(Ustr::from("BTC-USD")),
1547            inst_type: None,
1548            inst_family: None,
1549            bar: None,
1550        };
1551        let topic = topic_from_websocket_arg(&arg);
1552        handler.subscriptions_state.mark_subscribe(&topic);
1553        handler.subscriptions_state.confirm_subscribe(&topic);
1554        handler.subscriptions_state.mark_unsubscribe(&topic);
1555
1556        let rejected_subscription =
1557            handler.handle_subscription_error(&arg, "60019", "Unsubscription failed");
1558
1559        assert!(!rejected_subscription);
1560        assert_eq!(handler.subscriptions_state.all_topics(), vec![topic]);
1561        assert!(
1562            handler
1563                .subscriptions_state
1564                .pending_subscribe_topics()
1565                .is_empty()
1566        );
1567        assert!(
1568            handler
1569                .subscriptions_state
1570                .pending_unsubscribe_topics()
1571                .is_empty()
1572        );
1573    }
1574
1575    #[rstest]
1576    fn test_subscription_arg_from_error_message_matches_mainnet_shape() {
1577        let msg = "Wrong URL or channel:books,instId:BTC-USDT-SWAP doesn't exist. Please use the \
1578                   correct URL, channel and parameters referring to API document.";
1579
1580        let arg = subscription_arg_from_error_message(msg).unwrap();
1581
1582        assert_eq!(arg.channel, OKXWsChannel::Books);
1583        assert_eq!(arg.inst_id, Some(Ustr::from("BTC-USDT-SWAP")));
1584        assert_eq!(arg.inst_type, None);
1585        assert_eq!(arg.inst_family, None);
1586        assert_eq!(arg.bar, None);
1587    }
1588
1589    #[rstest]
1590    #[case::book_arg(
1591        Some(OKXWsChannel::Books),
1592        "60018",
1593        "Subscription failed",
1594        log::Level::Warn
1595    )]
1596    #[case::book_from_message(
1597        None,
1598        "60018",
1599        "Wrong URL or channel:books,instId:BTC-USDT-SWAP doesn't exist.",
1600        log::Level::Warn
1601    )]
1602    #[case::retryable(
1603        Some(OKXWsChannel::Tickers),
1604        "50001",
1605        "Service unavailable",
1606        log::Level::Warn
1607    )]
1608    #[case::non_book_arg(
1609        Some(OKXWsChannel::Tickers),
1610        "60018",
1611        "Subscription failed",
1612        log::Level::Error
1613    )]
1614    #[case::non_book_from_message(
1615        None,
1616        "60018",
1617        "Wrong URL or channel:tickers,instId:BTC-USDT-SWAP doesn't exist.",
1618        log::Level::Error
1619    )]
1620    fn test_error_frame_level(
1621        #[case] channel: Option<OKXWsChannel>,
1622        #[case] code: &str,
1623        #[case] msg: &str,
1624        #[case] expected: log::Level,
1625    ) {
1626        let arg = channel.map(|channel| OKXWebSocketArg {
1627            channel,
1628            inst_id: Some(Ustr::from("BTC-USDT-SWAP")),
1629            inst_type: None,
1630            inst_family: None,
1631            bar: None,
1632        });
1633
1634        let level = error_frame_level(arg.as_ref(), code, msg);
1635
1636        assert_eq!(level, expected);
1637    }
1638
1639    #[rstest]
1640    fn test_subscription_error_ignores_non_pending_topic() {
1641        let handler = create_handler();
1642        let arg = OKXWebSocketArg {
1643            channel: OKXWsChannel::Books,
1644            inst_id: Some(Ustr::from("BTC-USDT-SWAP")),
1645            inst_type: None,
1646            inst_family: None,
1647            bar: None,
1648        };
1649
1650        let rejected_subscription =
1651            handler.handle_subscription_error(&arg, "60018", "Subscription failed");
1652
1653        assert!(!rejected_subscription);
1654        assert!(handler.subscriptions_state.all_topics().is_empty());
1655        assert!(
1656            handler
1657                .subscriptions_state
1658                .pending_subscribe_topics()
1659                .is_empty()
1660        );
1661        assert!(
1662            handler
1663                .subscriptions_state
1664                .pending_unsubscribe_topics()
1665                .is_empty()
1666        );
1667    }
1668
1669    #[derive(serde::Deserialize, Debug, PartialEq)]
1670    struct ParseArrayItem {
1671        value: i64,
1672    }
1673
1674    #[rstest]
1675    fn test_parse_array_items_keeps_good_items_when_one_fails() {
1676        let data = json!([
1677            {"value": 1},
1678            {"value": "not a number"},
1679            {"value": 3},
1680        ]);
1681
1682        let parsed: Vec<ParseArrayItem> =
1683            parse_array_items(data, "test", false).expect("non-empty");
1684        assert_eq!(
1685            parsed,
1686            vec![ParseArrayItem { value: 1 }, ParseArrayItem { value: 3 }],
1687        );
1688    }
1689
1690    #[rstest]
1691    fn test_parse_array_items_returns_none_when_payload_not_array() {
1692        let data = json!({"not": "an array"});
1693        let parsed: Option<Vec<ParseArrayItem>> = parse_array_items(data, "test", false);
1694        assert!(parsed.is_none());
1695    }
1696
1697    #[rstest]
1698    fn test_parse_array_items_returns_none_when_all_items_fail() {
1699        let data = json!([{"value": "bad"}]);
1700        let parsed: Option<Vec<ParseArrayItem>> = parse_array_items(data, "test", false);
1701        assert!(parsed.is_none());
1702    }
1703
1704    #[rstest]
1705    fn test_route_instruments_keeps_valid_items_when_one_item_fails() {
1706        let handler = create_handler();
1707        let mut frame: Value =
1708            serde_json::from_str(&load_test_json("ws_instruments.json")).expect("valid fixture");
1709        let data = frame
1710            .get_mut("data")
1711            .and_then(Value::as_array_mut)
1712            .expect("data array");
1713        let mut invalid_item = data[0].clone();
1714        invalid_item["tickSz"] = json!(7);
1715        data.insert(0, invalid_item);
1716
1717        let arg: OKXWebSocketArg = serde_json::from_value(frame["arg"].clone()).expect("valid arg");
1718        let msg = handler
1719            .route_data_message(arg, frame["data"].clone())
1720            .expect("instruments message");
1721
1722        match msg {
1723            OKXWsMessage::Instruments(instruments) => {
1724                assert_eq!(instruments.len(), 1);
1725                assert_eq!(instruments[0].inst_id, "BTC-USDT-SWAP");
1726            }
1727            other => panic!("Expected Instruments, was {other:?}"),
1728        }
1729    }
1730
1731    #[rstest]
1732    fn test_route_instruments_prefers_rpi_over_legacy_alias() {
1733        let handler = create_handler();
1734        let mut frame: Value =
1735            serde_json::from_str(&load_test_json("ws_instruments.json")).expect("valid fixture");
1736        let instrument = &mut frame["data"][0];
1737        instrument["rpi"] = json!("2");
1738        instrument["elp"] = json!("1");
1739
1740        let arg: OKXWebSocketArg = serde_json::from_value(frame["arg"].clone()).expect("valid arg");
1741        let msg = handler
1742            .route_data_message(arg, frame["data"].clone())
1743            .expect("instruments message");
1744
1745        match msg {
1746            OKXWsMessage::Instruments(instruments) => {
1747                assert_eq!(instruments.len(), 1);
1748                assert_eq!(instruments[0].rpi, Some(OKXRpiPermission::Permitted));
1749            }
1750            other => panic!("Expected Instruments, was {other:?}"),
1751        }
1752    }
1753
1754    #[rstest]
1755    fn test_route_liquidation_warnings() {
1756        let handler = create_handler();
1757        let frame: Value = serde_json::from_str(&load_test_json("ws_liquidation_warning.json"))
1758            .expect("valid fixture");
1759
1760        let arg: OKXWebSocketArg = serde_json::from_value(frame["arg"].clone()).expect("valid arg");
1761        let msg = handler
1762            .route_data_message(arg, frame["data"].clone())
1763            .expect("liquidation warning message");
1764
1765        match msg {
1766            OKXWsMessage::LiquidationWarnings(warnings) => {
1767                assert_eq!(warnings.len(), 1);
1768                assert_eq!(warnings[0].inst_id, "BTC-USDT-SWAP");
1769                assert_eq!(warnings[0].mgn_ratio, "0.62");
1770            }
1771            other => panic!("Expected LiquidationWarnings, was {other:?}"),
1772        }
1773    }
1774}
nautilus_okx/websocket/handler.rs

nautilus_okx/websocket/
handler.rs
