Skip to main content

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        // Capture ownership before checking reconnect so recovery cannot select a replacement
240        let connection_epoch = client.connection_epoch();
241        if client.connection_mode().is_reconnect() {
242            return Err(OKXWsError::TransportSend(SendError::ConnectionChanged));
243        }
244
245        client
246            .send_text_on_connection(payload, rate_limit_keys, connection_epoch)
247            .await
248            .map_err(OKXWsError::TransportSend)
249    }
250
251    pub(super) async fn send_pong(&self) -> anyhow::Result<()> {
252        match self.send_on_connection(TEXT_PONG.to_string(), None).await {
253            Ok(()) => {
254                log::trace!("Sent pong response to OKX text ping");
255                Ok(())
256            }
257            Err(e) => Err(anyhow::anyhow!("Failed to send pong: {e}")),
258        }
259    }
260
261    pub(super) async fn next(&mut self) -> Option<OKXWsMessage> {
262        if let Some(message) = self.pending_messages.pop_front() {
263            return Some(message);
264        }
265
266        let mut poll_raw_next = false;
267
268        loop {
269            if self.signal.load(Ordering::Acquire) {
270                log::debug!("Stop signal received");
271                return None;
272            }
273
274            tokio::select! {
275                biased;
276                Some(cmd) = self.cmd_rx.recv(), if !poll_raw_next => {
277                    match cmd {
278                        HandlerCommand::SetClient(client) => {
279                            log::debug!("Handler received WebSocket client");
280                            self.inner = Some(client);
281                        }
282                        HandlerCommand::Disconnect => {
283                            log::debug!("Handler disconnecting WebSocket client");
284                            self.inner = None;
285                            return None;
286                        }
287                        HandlerCommand::Authenticate { payload } => {
288                            if let Err(e) = self.send_secret_with_retry(
289                                payload,
290                                Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()),
291                            ).await {
292                                log::error!(
293                                    "Failed to send authentication message after retries: error={e}"
294                                );
295                            }
296                        }
297                        HandlerCommand::Subscribe { args } => {
298                            if let Err(e) = self.handle_subscribe(args).await {
299                                log::error!("Failed to handle subscribe command: error={e}");
300                            }
301                        }
302                        HandlerCommand::SubscribeBook { subscription, cancel, gate, completion } => {
303                            let result = tokio::select! {
304                                biased;
305                                () = cancel.cancelled() => continue,
306                                result = async {
307                                    let client = self.inner.as_ref().ok_or(OKXWsError::NoActiveClient)?;
308                                    self.send_book_subscription(
309                                        OKXWsOperation::Subscribe,
310                                        subscription,
311                                        client.connection_epoch(),
312                                    ).await
313                                } => result,
314                            };
315
316                            if result.is_ok() {
317                                gate.open();
318                            }
319                            let _ = completion.send(result);
320                        }
321                        HandlerCommand::Unsubscribe { args } => {
322                            if let Err(e) = self.handle_unsubscribe(args).await {
323                                log::error!("Failed to handle unsubscribe command: error={e}");
324                            }
325                        }
326                        HandlerCommand::Resubscribe { subscription, cancel, gate, completion } => {
327                            let result = tokio::select! {
328                                biased;
329                                () = cancel.cancelled() => continue,
330                                result = async {
331                                    self.subscriptions_state.mark_failure(&topic_from_subscription_arg(&subscription));
332                                    let client = self.inner.as_ref().ok_or(OKXWsError::NoActiveClient)?;
333                                    let epoch = client.connection_epoch();
334                                    self.send_book_subscription(
335                                        OKXWsOperation::Unsubscribe,
336                                        subscription.clone(),
337                                        epoch,
338                                    ).await?;
339                                    self.send_book_subscription(
340                                        OKXWsOperation::Subscribe,
341                                        subscription,
342                                        epoch,
343                                    ).await
344                                } => result,
345                            };
346
347                            if result.is_ok() {
348                                gate.open();
349                            }
350                            let _ = completion.send(result);
351                        }
352                        HandlerCommand::Send {
353                            payload,
354                            rate_limit_keys,
355                            request_id,
356                            client_order_ids,
357                            op,
358                        } => {
359                            if let Err(e) = self.send_on_connection(
360                                payload,
361                                rate_limit_keys.as_deref(),
362                            ).await {
363                                log::debug!("Failed to send message: error={e}");
364
365                                if let Some(request_id) = request_id {
366                                    let message = OKXWsMessage::SendFailed {
367                                        request_id,
368                                        client_order_ids,
369                                        op,
370                                        error: e,
371                                    };
372
373                                    if self.send(message).is_err() {
374                                        if !self.is_stopped() {
375                                            log::error!("Failed to forward WebSocket send failure: receiver dropped");
376                                        }
377                                        return None;
378                                    }
379                                }
380                            }
381                        }
382                    }
383
384                    poll_raw_next = true;
385                }
386
387                msg = self.raw_rx.recv() => {
388                    let event = match msg {
389                        Some(msg) => match Self::parse_raw_message(msg) {
390                            Some(event) => event,
391                            None => continue,
392                        },
393                        None => {
394                            log::debug!("WebSocket stream closed");
395                            return None;
396                        }
397                    };
398
399                    match event {
400                        OKXWsFrame::Ping => {
401                            if let Err(e) = self.send_pong().await {
402                                log::warn!("Failed to send pong response: error={e}");
403                            }
404                        }
405                        OKXWsFrame::Login {
406                            code, msg, conn_id, ..
407                        } => {
408                            if code == OKX_SUCCESS_CODE {
409                                self.auth_tracker.succeed();
410                                return Some(OKXWsMessage::Authenticated);
411                            }
412
413                            self.auth_tracker.fail(msg.clone());
414
415                            let error = OKXWebSocketError {
416                                code,
417                                message: msg,
418                                conn_id: Some(conn_id),
419                                timestamp: self.clock.get_time_ns().as_u64(),
420                            };
421                            self.pending_messages.push_back(OKXWsMessage::Error(error));
422                        }
423                        OKXWsFrame::BookData { arg, action, data } => {
424                            return Some(OKXWsMessage::BookData { arg, action, data });
425                        }
426                        OKXWsFrame::RpiBookData { arg, action, data } => {
427                            return Some(OKXWsMessage::RpiBookData { arg, action, data });
428                        }
429                        OKXWsFrame::OrderResponse {
430                            id, op, code, msg, data,
431                        } => {
432                            return Some(OKXWsMessage::OrderResponse {
433                                id, op, code, msg, data,
434                            });
435                        }
436                        OKXWsFrame::Data { arg, data } => {
437                            if let Some(output) = self.route_data_message(arg, data) {
438                                return Some(output);
439                            }
440                        }
441                        OKXWsFrame::Error { arg, code, msg } => {
442                            let arg = arg.or_else(|| subscription_arg_from_error_message(&msg));
443                            if let Some(arg) = arg
444                                && self.handle_subscription_error(&arg, &code, &msg)
445                            {
446                                return Some(OKXWsMessage::SubscriptionFailed {
447                                    channel: arg.channel,
448                                    inst_id: arg.inst_id,
449                                    code,
450                                    msg,
451                                });
452                            }
453
454                            let error = OKXWebSocketError {
455                                code,
456                                message: msg,
457                                conn_id: None,
458                                timestamp: self.clock.get_time_ns().as_u64(),
459                            };
460                            return Some(OKXWsMessage::Error(error));
461                        }
462                        OKXWsFrame::Reconnected => {
463                            self.auth_tracker.invalidate();
464                            return Some(OKXWsMessage::Reconnected);
465                        }
466                        OKXWsFrame::Subscription {
467                            event, arg, code, msg,
468                            ..
469                        } => {
470                            let rejected = self
471                                .handle_subscription_ack(&event, &arg, code.as_deref(), msg.as_deref());
472
473                            if rejected {
474                                return Some(OKXWsMessage::SubscriptionFailed {
475                                    channel: arg.channel,
476                                    inst_id: arg.inst_id,
477                                    code: code.unwrap_or_default(),
478                                    msg: msg.unwrap_or_default(),
479                                });
480                            }
481                        }
482                        OKXWsFrame::ChannelConnCount { .. } => {}
483                    }
484                }
485
486                () = std::future::ready(()), if poll_raw_next => {
487                    poll_raw_next = false;
488                }
489
490                // Polled last so a ready message never registers and drops a timer
491                () = time::sleep(time::Duration::from_millis(100)) => {
492                    // Wake the loop to poll the stop signal while both channels are idle
493                }
494
495                else => {
496                    log::debug!("Handler shutting down: stream ended or command channel closed");
497                    return None;
498                }
499            }
500        }
501    }
502
503    fn route_data_message(&self, arg: OKXWebSocketArg, mut data: Value) -> Option<OKXWsMessage> {
504        let OKXWebSocketArg {
505            channel, inst_id, ..
506        } = arg;
507
508        match channel {
509            OKXWsChannel::Account => Some(OKXWsMessage::Account(data)),
510            OKXWsChannel::Positions => Some(OKXWsMessage::Positions(data)),
511            OKXWsChannel::Orders => {
512                parse_array_items(data, "orders", false).map(OKXWsMessage::Orders)
513            }
514            OKXWsChannel::SprdOrders => {
515                parse_array_items(data, "spread orders", false).map(OKXWsMessage::SpreadOrders)
516            }
517            OKXWsChannel::OrdersAlgo | OKXWsChannel::AlgoAdvance => {
518                parse_array_items(data, "algo orders", false).map(OKXWsMessage::AlgoOrders)
519            }
520            OKXWsChannel::LiquidationWarning => {
521                parse_array_items(data, "liquidation warnings", false)
522                    .map(OKXWsMessage::LiquidationWarnings)
523            }
524            OKXWsChannel::Instruments => {
525                prefer_rpi_response_fields(&mut data);
526                parse_array_items(data, "instruments", true).map(OKXWsMessage::Instruments)
527            }
528            _ => Some(OKXWsMessage::ChannelData {
529                channel,
530                inst_id,
531                data,
532            }),
533        }
534    }
535
536    fn handle_subscription_ack(
537        &self,
538        event: &OKXSubscriptionEvent,
539        arg: &OKXWebSocketArg,
540        code: Option<&str>,
541        msg: Option<&str>,
542    ) -> bool {
543        let topic = topic_from_websocket_arg(arg);
544        let success = code.is_none_or(|c| c == OKX_SUCCESS_CODE);
545
546        match event {
547            OKXSubscriptionEvent::Subscribe => {
548                if success {
549                    self.subscriptions_state.confirm_subscribe(&topic);
550                    false
551                } else {
552                    log::warn!(
553                        "Subscription failed: topic={topic:?}, error={msg:?}, code={code:?}"
554                    );
555                    self.subscriptions_state.mark_failure(&topic);
556                    true
557                }
558            }
559            OKXSubscriptionEvent::Unsubscribe => {
560                if success {
561                    self.subscriptions_state.confirm_unsubscribe(&topic);
562                } else {
563                    log::warn!(
564                        "Unsubscription failed - restoring subscription: \
565                         topic={topic:?}, error={msg:?}, code={code:?}"
566                    );
567                    self.subscriptions_state.confirm_unsubscribe(&topic);
568                    self.subscriptions_state.mark_subscribe(&topic);
569                    self.subscriptions_state.confirm_subscribe(&topic);
570                }
571                false
572            }
573        }
574    }
575
576    fn handle_subscription_error(&self, arg: &OKXWebSocketArg, code: &str, msg: &str) -> bool {
577        let topic = topic_from_websocket_arg(arg);
578        let event = if self
579            .subscriptions_state
580            .pending_unsubscribe_topics()
581            .iter()
582            .any(|pending| pending == &topic)
583        {
584            OKXSubscriptionEvent::Unsubscribe
585        } else if self
586            .subscriptions_state
587            .pending_subscribe_topics()
588            .iter()
589            .any(|pending| pending == &topic)
590            || (arg.channel.is_book()
591                && self
592                    .subscriptions_state
593                    .all_topics()
594                    .iter()
595                    .any(|tracked| tracked == &topic))
596        {
597            OKXSubscriptionEvent::Subscribe
598        } else {
599            return false;
600        };
601
602        self.handle_subscription_ack(&event, arg, Some(code), Some(msg))
603    }
604
605    async fn send_book_subscription(
606        &self,
607        op: OKXWsOperation,
608        subscription: OKXSubscriptionArg,
609        connection_epoch: u64,
610    ) -> Result<(), OKXWsError> {
611        let client = self.inner.as_ref().ok_or(OKXWsError::NoActiveClient)?;
612        let message = OKXSubscription {
613            op,
614            args: vec![subscription],
615        };
616        let payload =
617            serde_json::to_string(&message).map_err(|e| OKXWsError::ClientError(e.to_string()))?;
618
619        client
620            .send_text_on_connection(
621                payload,
622                Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()),
623                connection_epoch,
624            )
625            .await
626            .map_err(OKXWsError::TransportSend)
627    }
628
629    async fn handle_subscribe(&self, args: Vec<OKXSubscriptionArg>) -> anyhow::Result<()> {
630        for arg in &args {
631            log::debug!(
632                "Subscribing to channel: channel={:?}, inst_id={:?}",
633                arg.channel,
634                arg.inst_id
635            );
636        }
637
638        let message = OKXSubscription {
639            op: OKXWsOperation::Subscribe,
640            args,
641        };
642
643        let json_txt = serde_json::to_string(&message)
644            .map_err(|e| anyhow::anyhow!("Failed to serialize subscription: {e}"))?;
645
646        self.send_with_retry(json_txt, Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
647            .await
648            .map_err(|e| anyhow::anyhow!("Failed to send subscription after retries: {e}"))?;
649        Ok(())
650    }
651
652    async fn handle_unsubscribe(&self, args: Vec<OKXSubscriptionArg>) -> anyhow::Result<()> {
653        for arg in &args {
654            log::debug!(
655                "Unsubscribing from channel: channel={:?}, inst_id={:?}",
656                arg.channel,
657                arg.inst_id
658            );
659        }
660
661        let message = OKXSubscription {
662            op: OKXWsOperation::Unsubscribe,
663            args,
664        };
665
666        let json_txt = serde_json::to_string(&message)
667            .map_err(|e| anyhow::anyhow!("Failed to serialize unsubscription: {e}"))?;
668
669        self.send_with_retry(json_txt, Some(OKX_RATE_LIMIT_KEY_SUBSCRIPTION.as_slice()))
670            .await
671            .map_err(|e| anyhow::anyhow!("Failed to send unsubscription after retries: {e}"))?;
672        Ok(())
673    }
674
675    pub(crate) fn parse_raw_message(
676        msg: tokio_tungstenite::tungstenite::Message,
677    ) -> Option<OKXWsFrame> {
678        match msg {
679            tokio_tungstenite::tungstenite::Message::Text(text) => {
680                if text == TEXT_PONG {
681                    log::trace!("Received pong from OKX");
682                    return None;
683                }
684
685                if text == TEXT_PING {
686                    log::trace!("Received ping from OKX (text)");
687                    return Some(OKXWsFrame::Ping);
688                }
689
690                if text == RECONNECTED {
691                    log::debug!("Received WebSocket reconnection signal");
692                    return Some(OKXWsFrame::Reconnected);
693                }
694                log::trace!("Received WebSocket message: {text}");
695
696                match serde_json::from_str(&text) {
697                    Ok(ws_event) => match &ws_event {
698                        OKXWsFrame::Error { arg, code, msg } => {
699                            let level = error_frame_level(arg.as_ref(), code, msg);
700                            log::log!(level, "WebSocket error: {code} - {msg}");
701                            Some(ws_event)
702                        }
703                        OKXWsFrame::Login {
704                            event,
705                            code,
706                            msg,
707                            conn_id,
708                        } => {
709                            if code == OKX_SUCCESS_CODE {
710                                log::debug!("WebSocket authenticated: conn_id={conn_id}");
711                            } else {
712                                log::error!(
713                                    "WebSocket authentication failed: \
714                                     event={event}, code={code}, error={msg}"
715                                );
716                            }
717                            Some(ws_event)
718                        }
719                        OKXWsFrame::Subscription {
720                            event,
721                            arg,
722                            conn_id,
723                            ..
724                        } => {
725                            let channel_str = serde_json::to_string(&arg.channel)
726                                .expect("Invalid OKX websocket channel")
727                                .trim_matches('"')
728                                .to_string();
729                            log::debug!("{event}d: channel={channel_str}, conn_id={conn_id}");
730                            Some(ws_event)
731                        }
732                        OKXWsFrame::ChannelConnCount {
733                            channel,
734                            conn_count,
735                            conn_id,
736                            ..
737                        } => {
738                            let channel_str = serde_json::to_string(channel)
739                                .expect("Invalid OKX websocket channel")
740                                .trim_matches('"')
741                                .to_string();
742                            log::debug!(
743                                "Channel connection status: \
744                                 channel={channel_str}, connections={conn_count}, conn_id={conn_id}",
745                            );
746                            None
747                        }
748                        OKXWsFrame::Ping => {
749                            log::trace!("Ignoring ping event parsed from text payload");
750                            None
751                        }
752                        OKXWsFrame::Data { .. }
753                        | OKXWsFrame::BookData { .. }
754                        | OKXWsFrame::RpiBookData { .. } => Some(ws_event),
755                        OKXWsFrame::OrderResponse {
756                            id, op, code, data, ..
757                        } => {
758                            if code == OKX_SUCCESS_CODE {
759                                log::debug!(
760                                    "Order operation successful: id={id:?}, op={op}, code={code}"
761                                );
762
763                                if let Some(order_data) = data.first() {
764                                    let success_msg = order_data
765                                        .get(OKX_FIELD_SMSG)
766                                        .and_then(|s| s.as_str())
767                                        .unwrap_or("Order operation successful");
768                                    log::debug!("Order success details: {success_msg}");
769                                }
770                            }
771                            Some(ws_event)
772                        }
773                        OKXWsFrame::Reconnected => {
774                            log::warn!("Unexpected Reconnected event from deserialization");
775                            None
776                        }
777                    },
778                    Err(e) => {
779                        log::error!("Failed to parse message: {e}: {text}");
780                        None
781                    }
782                }
783            }
784            Message::Ping(_payload) => {
785                log::trace!("Received binary ping frame from OKX");
786                Some(OKXWsFrame::Ping)
787            }
788            Message::Pong(payload) => {
789                log::trace!("Received pong frame from OKX ({} bytes)", payload.len());
790                None
791            }
792            Message::Binary(msg) => {
793                log::debug!("Raw binary frame ({} bytes)", msg.len());
794                log::trace!("Raw binary: {msg:?}");
795                None
796            }
797            Message::Close(_) => {
798                log::debug!("Received close message");
799                None
800            }
801            msg => {
802                log::warn!("Unexpected message: {msg}");
803                None
804            }
805        }
806    }
807}
808
809// Book recovery retries a rejected book subscription and reports its own failure
810fn error_frame_level(arg: Option<&OKXWebSocketArg>, code: &str, msg: &str) -> log::Level {
811    let book = match arg {
812        Some(arg) => arg.channel.is_book(),
813        None => subscription_arg_from_error_message(msg).is_some_and(|arg| arg.channel.is_book()),
814    };
815
816    if should_retry_error_code(code) || book {
817        log::Level::Warn
818    } else {
819        log::Level::Error
820    }
821}
822
823fn subscription_arg_from_error_message(msg: &str) -> Option<OKXWebSocketArg> {
824    let descriptor = msg
825        .strip_prefix("Wrong URL or channel:")?
826        .split_whitespace()
827        .next()?;
828    let mut fields = descriptor.split(',');
829    let channel = fields.next()?;
830    let mut arg = Map::new();
831    arg.insert("channel".to_string(), Value::String(channel.to_string()));
832
833    for field in fields {
834        let (key, value) = field.split_once(':')?;
835        if !matches!(key, "instId" | "sprdId" | "instType" | "instFamily") {
836            return None;
837        }
838        arg.insert(key.to_string(), Value::String(value.to_string()));
839    }
840
841    serde_json::from_value(Value::Object(arg)).ok()
842}
843
844/// Returns `true` when an OKX WebSocket order message represents a post-only auto-cancel.
845pub fn is_post_only_auto_cancel(msg: &OKXOrderMsg) -> bool {
846    use crate::common::{consts::OKX_POST_ONLY_CANCEL_SOURCE, enums::OKXOrderStatus};
847
848    if msg.state != OKXOrderStatus::Canceled {
849        return false;
850    }
851
852    let cancel_source_matches = matches!(
853        msg.cancel_source.as_deref(),
854        Some(source) if source == OKX_POST_ONLY_CANCEL_SOURCE
855    );
856
857    let reason_matches = matches!(
858        msg.cancel_source_reason.as_deref(),
859        Some(reason) if reason.contains("POST_ONLY")
860    );
861
862    if !(cancel_source_matches || reason_matches) {
863        return false;
864    }
865
866    msg.acc_fill_sz
867        .as_ref()
868        .is_none_or(|filled| filled == "0" || filled.is_empty())
869}
870
871/// Returns `true` when an RPI order update is canceled without any fill.
872pub fn is_unfilled_rpi_cancel(msg: &OKXOrderMsg) -> bool {
873    msg.ord_type == OKXOrderType::Rpi
874        && msg.state == OKXOrderStatus::Canceled
875        && msg
876            .acc_fill_sz
877            .as_ref()
878            .is_none_or(|filled| filled == "0" || filled.is_empty())
879}
880
881// Per-item deserialization so one malformed entry does not drop the batch.
882fn parse_array_items<T: serde::de::DeserializeOwned>(
883    data: Value,
884    label: &str,
885    warn_on_parse_error: bool,
886) -> Option<Vec<T>> {
887    let Value::Array(items) = data else {
888        if warn_on_parse_error {
889            log::warn!("Expected {label} payload to be a JSON array");
890        } else {
891            log::error!("Expected {label} payload to be a JSON array");
892        }
893        return None;
894    };
895
896    let mut parsed = Vec::with_capacity(items.len());
897    for (idx, item) in items.into_iter().enumerate() {
898        match serde_json::from_value::<T>(item) {
899            Ok(value) => parsed.push(value),
900            Err(e) => {
901                if warn_on_parse_error {
902                    log::warn!("Failed to parse {label} item at index {idx}: {e}");
903                } else {
904                    log::error!("Failed to parse {label} item at index {idx}: {e}");
905                }
906            }
907        }
908    }
909
910    if parsed.is_empty() {
911        None
912    } else {
913        Some(parsed)
914    }
915}
916
917fn should_retry_replay_safe_error(error: &OKXWsError) -> bool {
918    match error {
919        OKXWsError::OkxError { error_code, .. } => should_retry_error_code(error_code),
920        OKXWsError::TransportSend(SendError::Timeout | SendError::ConnectionChanged)
921        | OKXWsError::TungsteniteError(_)
922        | OKXWsError::OperationTimeout { .. } => true,
923        OKXWsError::AuthenticationError(_)
924        | OKXWsError::JsonError(_)
925        | OKXWsError::ParsingError(_)
926        | OKXWsError::ClientError(_)
927        | OKXWsError::NoActiveClient
928        | OKXWsError::HandlerUnavailable(_)
929        | OKXWsError::TransportSend(
930            SendError::InvalidInput(_)
931            | SendError::BufferFull
932            | SendError::Closed
933            | SendError::WriteTimeout
934            | SendError::BrokenPipe(_),
935        )
936        | OKXWsError::SendFailed(_) => false,
937    }
938}
939
940fn create_okx_retry_error(error: RetryError) -> OKXWsError {
941    match error {
942        RetryError::OperationTimeout { timeout_ms } => OKXWsError::OperationTimeout { timeout_ms },
943        RetryError::InvalidConfiguration { message } => OKXWsError::ClientError(message),
944        RetryError::Canceled => {
945            OKXWsError::SendFailed("Adapter disconnecting or shutting down".to_string())
946        }
947        error @ RetryError::ElapsedBudgetExceeded { .. } => {
948            OKXWsError::SendFailed(error.to_string())
949        }
950    }
951}
952
953#[cfg(test)]
954mod tests {
955    use std::{
956        num::NonZeroU32,
957        sync::{Arc, atomic::AtomicBool},
958        time::Duration,
959    };
960
961    use futures_util::{SinkExt, StreamExt};
962    use nautilus_core::{collections::AtomicMap, time::get_atomic_clock_realtime};
963    use nautilus_model::identifiers::InstrumentId;
964    use nautilus_network::{
965        mode::ConnectionMode,
966        ratelimiter::{RateLimiter, quota::Quota},
967        websocket::{AuthTracker, SubscriptionState, WebSocketConfig, channel_message_handler},
968    };
969    use rstest::rstest;
970    use serde_json::json;
971
972    use super::*;
973    use crate::{
974        book::{BookChannelScope, BookSequenceOutcome, sync::BookSyncTracker},
975        common::{
976            consts::OKX_WS_TOPIC_DELIMITER,
977            enums::{OKXBookAction, OKXBookChannel, OKXRpiPermission},
978            testing::load_test_json,
979        },
980    };
981
982    fn create_handler() -> OKXWsFeedHandler {
983        let signal = Arc::new(AtomicBool::new(false));
984        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
985        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
986        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
987
988        OKXWsFeedHandler::new(
989            signal,
990            cmd_rx,
991            raw_rx,
992            out_tx,
993            AuthTracker::new(),
994            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
995            get_atomic_clock_realtime(),
996        )
997    }
998
999    #[rstest]
1000    fn test_command_debug_redacts_payloads() {
1001        let payload = "authentication-secret";
1002        let authenticate = HandlerCommand::Authenticate {
1003            payload: SecretString::from(payload.to_string()),
1004        };
1005        let send = HandlerCommand::Send {
1006            payload: payload.to_string(),
1007            rate_limit_keys: None,
1008            request_id: None,
1009            client_order_ids: Vec::new(),
1010            op: None,
1011        };
1012
1013        let debug = format!("{authenticate:?} {send:?}");
1014
1015        assert!(debug.contains(REDACTED));
1016        assert!(!debug.contains(payload));
1017    }
1018
1019    #[rstest]
1020    #[case::submit(OKXWsOperation::Order)]
1021    #[case::cancel(OKXWsOperation::CancelOrder)]
1022    #[case::amend(OKXWsOperation::AmendOrder)]
1023    #[tokio::test]
1024    async fn send_failure_returns_without_inbound_frame(#[case] op: OKXWsOperation) {
1025        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1026        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1027        let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
1028
1029        let mut handler = OKXWsFeedHandler::new(
1030            Arc::new(AtomicBool::new(false)),
1031            cmd_rx,
1032            raw_rx,
1033            out_tx,
1034            AuthTracker::new(),
1035            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1036            get_atomic_clock_realtime(),
1037        );
1038        let client_order_id = ClientOrderId::from("O-disconnected");
1039        cmd_tx
1040            .send(HandlerCommand::Send {
1041                payload: String::new(),
1042                rate_limit_keys: None,
1043                request_id: Some("disconnected-request".to_string()),
1044                client_order_ids: vec![client_order_id],
1045                op: Some(op.clone()),
1046            })
1047            .unwrap();
1048
1049        let message = tokio::time::timeout(Duration::from_secs(1), async {
1050            tokio::select! {
1051                message = out_rx.recv() => message,
1052                message = handler.next() => panic!("Unexpected handler return: {message:?}"),
1053            }
1054        })
1055        .await
1056        .expect("send failure must not wait for an inbound frame");
1057
1058        let Some(OKXWsMessage::SendFailed {
1059            request_id,
1060            client_order_ids,
1061            op: failed_op,
1062            error,
1063        }) = message
1064        else {
1065            panic!("Expected send failure, received {message:?}");
1066        };
1067
1068        assert_eq!(request_id, "disconnected-request");
1069        assert_eq!(client_order_ids, vec![client_order_id]);
1070        assert_eq!(failed_op, Some(op));
1071        assert!(matches!(error, OKXWsError::NoActiveClient));
1072    }
1073
1074    #[tokio::test]
1075    async fn send_failures_preserve_raw_message_fairness() {
1076        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1077        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1078        let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
1079
1080        let mut handler = OKXWsFeedHandler::new(
1081            Arc::new(AtomicBool::new(false)),
1082            cmd_rx,
1083            raw_rx,
1084            out_tx,
1085            AuthTracker::new(),
1086            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1087            get_atomic_clock_realtime(),
1088        );
1089
1090        for request_id in ["first-cancel", "second-cancel"] {
1091            cmd_tx
1092                .send(HandlerCommand::Send {
1093                    payload: String::new(),
1094                    rate_limit_keys: None,
1095                    request_id: Some(request_id.to_string()),
1096                    client_order_ids: vec![ClientOrderId::from("O-disconnected")],
1097                    op: Some(OKXWsOperation::CancelOrder),
1098                })
1099                .unwrap();
1100        }
1101
1102        raw_tx
1103            .send(Message::Text(RECONNECTED.to_string().into()))
1104            .unwrap();
1105
1106        raw_tx
1107            .send(Message::Text(RECONNECTED.to_string().into()))
1108            .unwrap();
1109
1110        let messages = tokio::time::timeout(Duration::from_secs(1), async {
1111            for _ in 0..2 {
1112                let message = handler.next().await.expect("raw message");
1113                handler.send(message).unwrap();
1114            }
1115
1116            [
1117                out_rx.recv().await,
1118                out_rx.recv().await,
1119                out_rx.recv().await,
1120                out_rx.recv().await,
1121            ]
1122        })
1123        .await
1124        .expect("failures and reconnect must arrive without more input");
1125
1126        assert!(
1127            matches!(&messages[0], Some(OKXWsMessage::SendFailed { request_id, .. })
1128            if request_id == "first-cancel")
1129        );
1130        assert!(matches!(messages[1], Some(OKXWsMessage::Reconnected)));
1131        assert!(
1132            matches!(&messages[2], Some(OKXWsMessage::SendFailed { request_id, .. })
1133            if request_id == "second-cancel")
1134        );
1135        assert!(matches!(messages[3], Some(OKXWsMessage::Reconnected)));
1136        assert!(handler.pending_messages.is_empty());
1137        assert!(handler.cmd_rx.is_empty());
1138    }
1139
1140    #[tokio::test]
1141    async fn test_next_polls_raw_after_one_ready_command() {
1142        let signal = Arc::new(AtomicBool::new(false));
1143        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1144        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1145        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1146        let mut handler = OKXWsFeedHandler::new(
1147            signal,
1148            cmd_rx,
1149            raw_rx,
1150            out_tx,
1151            AuthTracker::new(),
1152            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1153            get_atomic_clock_realtime(),
1154        );
1155
1156        for _ in 0..3 {
1157            cmd_tx
1158                .send(HandlerCommand::Subscribe { args: Vec::new() })
1159                .unwrap();
1160        }
1161        raw_tx
1162            .send(Message::Text(RECONNECTED.to_string().into()))
1163            .unwrap();
1164
1165        let message = handler.next().await;
1166
1167        assert!(matches!(message, Some(OKXWsMessage::Reconnected)));
1168        assert_eq!(handler.cmd_rx.len(), 2);
1169    }
1170
1171    #[rstest]
1172    #[case::sent(false)]
1173    #[case::canceled(true)]
1174    #[tokio::test]
1175    async fn initial_book_completion_waits_for_send(#[case] canceled: bool) {
1176        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1177        let url = format!("ws://{}/", listener.local_addr().unwrap());
1178        let (wire_tx, mut wire_rx) = tokio::sync::mpsc::unbounded_channel();
1179
1180        let server = tokio::spawn(async move {
1181            let (socket, _) = listener.accept().await.unwrap();
1182            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1183            while let Some(Ok(frame)) = socket.next().await {
1184                wire_tx.send(frame).unwrap();
1185            }
1186        });
1187
1188        let (message_handler, raw_rx) = channel_message_handler();
1189        let client = WebSocketClient::builder()
1190            .config(WebSocketConfig::builder().url(url).build().unwrap())
1191            .message_handler(message_handler)
1192            .default_quota(Quota::with_period(Duration::from_secs(10)).unwrap())
1193            .connect()
1194            .await
1195            .unwrap();
1196        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1197        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1198
1199        let mut handler = OKXWsFeedHandler::new(
1200            Arc::new(AtomicBool::new(false)),
1201            cmd_rx,
1202            raw_rx,
1203            out_tx,
1204            AuthTracker::new(),
1205            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1206            get_atomic_clock_realtime(),
1207        );
1208        handler.inner = Some(client);
1209        cmd_tx
1210            .send(HandlerCommand::Subscribe { args: Vec::new() })
1211            .unwrap();
1212        let cancel = CancellationToken::new();
1213        let gate = SnapshotGate::default();
1214        gate.lock().close();
1215        let (completion, mut result) = tokio::sync::oneshot::channel();
1216        cmd_tx
1217            .send(HandlerCommand::SubscribeBook {
1218                subscription: OKXSubscriptionArg {
1219                    channel: OKXWsChannel::Books,
1220                    inst_id: Some(Ustr::from("BTC-USDT")),
1221                    inst_type: None,
1222                    inst_family: None,
1223                },
1224                cancel: cancel.clone(),
1225                gate: gate.clone(),
1226                completion,
1227            })
1228            .unwrap();
1229
1230        let running = tokio::spawn(async move { handler.next().await });
1231        let first = tokio::time::timeout(Duration::from_secs(3), wire_rx.recv())
1232            .await
1233            .unwrap()
1234            .unwrap();
1235        let first: Value = serde_json::from_str(first.to_text().unwrap()).unwrap();
1236        assert_eq!(first, serde_json::json!({"op": "subscribe", "args": []}));
1237        assert!(matches!(
1238            result.try_recv(),
1239            Err(tokio::sync::oneshot::error::TryRecvError::Empty)
1240        ));
1241
1242        if canceled {
1243            cancel.cancel();
1244        }
1245
1246        tokio::time::pause();
1247        tokio::time::advance(Duration::from_secs(10)).await;
1248        tokio::time::resume();
1249        let result = tokio::time::timeout(Duration::from_secs(3), result)
1250            .await
1251            .unwrap();
1252
1253        if canceled {
1254            assert!(gate.lock().is_closed());
1255            assert!(result.is_err());
1256            assert!(
1257                tokio::time::timeout(Duration::from_millis(100), wire_rx.recv())
1258                    .await
1259                    .is_err()
1260            );
1261        } else {
1262            result.unwrap().unwrap();
1263            assert!(!gate.lock().is_closed());
1264            let frame = tokio::time::timeout(Duration::from_secs(3), wire_rx.recv())
1265                .await
1266                .unwrap()
1267                .unwrap();
1268            let frame: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
1269            assert_eq!(
1270                frame,
1271                serde_json::json!({"op": "subscribe", "args": [{"channel": "books", "instId": "BTC-USDT"}]})
1272            );
1273        }
1274
1275        running.abort();
1276        server.abort();
1277    }
1278
1279    #[tokio::test]
1280    async fn book_subscription_rejects_changed_connection() {
1281        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1282        let url = format!("ws://{}/", listener.local_addr().unwrap());
1283        let (wire_tx, mut wire_rx) = tokio::sync::mpsc::unbounded_channel();
1284
1285        let server = tokio::spawn(async move {
1286            let (socket, _) = listener.accept().await.unwrap();
1287            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1288            while let Some(Ok(frame)) = socket.next().await {
1289                wire_tx.send(frame).unwrap();
1290            }
1291        });
1292        let (message_handler, raw_rx) = channel_message_handler();
1293        let client = WebSocketClient::builder()
1294            .config(WebSocketConfig::builder().url(url).build().unwrap())
1295            .message_handler(message_handler)
1296            .connect()
1297            .await
1298            .unwrap();
1299        let wrong_epoch = client.connection_epoch() + 1;
1300        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1301        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1302        let mut handler = OKXWsFeedHandler::new(
1303            Arc::new(AtomicBool::new(false)),
1304            cmd_rx,
1305            raw_rx,
1306            out_tx,
1307            AuthTracker::new(),
1308            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1309            get_atomic_clock_realtime(),
1310        );
1311        handler.inner = Some(client);
1312        let result = handler
1313            .send_book_subscription(
1314                OKXWsOperation::Subscribe,
1315                OKXSubscriptionArg {
1316                    channel: OKXWsChannel::Books,
1317                    inst_id: Some(Ustr::from("BTC-USDT")),
1318                    inst_type: None,
1319                    inst_family: None,
1320                },
1321                wrong_epoch,
1322            )
1323            .await;
1324
1325        assert!(matches!(
1326            result,
1327            Err(OKXWsError::TransportSend(SendError::ConnectionChanged))
1328        ));
1329        assert!(matches!(
1330            wire_rx.try_recv(),
1331            Err(tokio::sync::mpsc::error::TryRecvError::Empty)
1332        ));
1333        handler.inner.as_ref().unwrap().disconnect().await;
1334        server.abort();
1335    }
1336
1337    #[tokio::test]
1338    async fn recovery_cancel_between_sends_keeps_snapshot_gate_closed() {
1339        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1340        let url = format!("ws://{}/", listener.local_addr().unwrap());
1341        let (received, first_frame) = tokio::sync::oneshot::channel();
1342
1343        let server = tokio::spawn(async move {
1344            let (socket, _) = listener.accept().await.unwrap();
1345            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1346            let frame = socket.next().await.unwrap().unwrap();
1347            received.send(frame).unwrap();
1348
1349            while socket.next().await.is_some() {}
1350        });
1351
1352        let (message_handler, raw_rx) = channel_message_handler();
1353        let client = WebSocketClient::builder()
1354            .config(WebSocketConfig::builder().url(url).build().unwrap())
1355            .message_handler(message_handler)
1356            .default_quota(Quota::per_hour(NonZeroU32::new(1).unwrap()))
1357            .connect()
1358            .await
1359            .unwrap();
1360        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1361        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1362
1363        let mut handler = OKXWsFeedHandler::new(
1364            Arc::new(AtomicBool::new(false)),
1365            cmd_rx,
1366            raw_rx,
1367            out_tx,
1368            AuthTracker::new(),
1369            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1370            get_atomic_clock_realtime(),
1371        );
1372        handler.inner = Some(client);
1373        let tracker = BookSyncTracker::default();
1374        let instrument_id = InstrumentId::from("BTC-USDT.OKX");
1375        tracker.record_subscription(instrument_id, time::Instant::now(), SnapshotGate::default());
1376        let recovery = tracker.claim_recovery(instrument_id).unwrap();
1377        assert!(recovery.begin_replacement());
1378        let cancel = recovery.cancellation.child_token();
1379        let (completion, result) = tokio::sync::oneshot::channel();
1380        cmd_tx
1381            .send(HandlerCommand::Resubscribe {
1382                subscription: OKXSubscriptionArg {
1383                    channel: OKXWsChannel::Books,
1384                    inst_id: Some(Ustr::from("BTC-USDT")),
1385                    inst_type: None,
1386                    inst_family: None,
1387                },
1388                cancel: cancel.clone(),
1389                gate: recovery.gate.clone(),
1390                completion,
1391            })
1392            .unwrap();
1393
1394        let running = tokio::spawn(async move { handler.next().await });
1395        let frame = tokio::time::timeout(Duration::from_secs(3), first_frame)
1396            .await
1397            .unwrap()
1398            .unwrap();
1399        let request: Value = serde_json::from_str(frame.to_text().unwrap()).unwrap();
1400        assert_eq!(request["op"], "unsubscribe");
1401        assert_eq!(
1402            tracker.validate_sequence(instrument_id, true, &[(Some(-1), 42)], time::Instant::now()),
1403            BookSequenceOutcome::Suppress
1404        );
1405        cancel.cancel();
1406        assert!(
1407            tokio::time::timeout(Duration::from_secs(1), result)
1408                .await
1409                .unwrap()
1410                .is_err()
1411        );
1412
1413        assert!(recovery.gate.lock().is_closed());
1414        assert_eq!(
1415            tracker.validate_sequence(instrument_id, true, &[(Some(-1), 43)], time::Instant::now()),
1416            BookSequenceOutcome::Suppress
1417        );
1418        cmd_tx.send(HandlerCommand::Disconnect).unwrap();
1419        tokio::time::timeout(Duration::from_secs(3), running)
1420            .await
1421            .unwrap()
1422            .unwrap();
1423        server.abort();
1424    }
1425
1426    #[rstest]
1427    #[case::disabled_deadline(Duration::ZERO)]
1428    #[case::snapshot_deadline(Duration::from_secs(3))]
1429    #[tokio::test]
1430    async fn reconnect_replay_survives_inflight_recovery(#[case] snapshot_timeout: Duration) {
1431        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1432        let url = format!("ws://{}/", listener.local_addr().unwrap());
1433        let (wire_tx, mut wire_rx) = tokio::sync::mpsc::unbounded_channel();
1434
1435        let server = tokio::spawn(async move {
1436            let (socket, _) = listener.accept().await.unwrap();
1437            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1438
1439            while let Some(Ok(Message::Text(text))) = socket.next().await {
1440                let request: Value = serde_json::from_str(&text).unwrap();
1441                wire_tx
1442                    .send(request["op"].as_str().unwrap().to_owned())
1443                    .unwrap();
1444
1445                if request["op"] == "subscribe" {
1446                    for (fixture, sequence, previous) in [
1447                        ("ws_books_snapshot.json", 100, -1),
1448                        ("ws_books_update.json", 101, 100),
1449                    ] {
1450                        let mut frame: Value =
1451                            serde_json::from_str(&load_test_json(fixture)).unwrap();
1452                        frame["data"][0]["seqId"] = json!(sequence);
1453                        frame["data"][0]["prevSeqId"] = json!(previous);
1454                        socket
1455                            .send(Message::Text(frame.to_string().into()))
1456                            .await
1457                            .unwrap();
1458                    }
1459                }
1460            }
1461        });
1462
1463        // Replay and unsubscribe consume the burst; the replacement subscribe cannot send
1464        let limiter = Arc::new(RateLimiter::new_with_quota(
1465            Some(
1466                Quota::with_period(Duration::from_secs(10))
1467                    .unwrap()
1468                    .allow_burst(NonZeroU32::new(2).unwrap()),
1469            ),
1470            Vec::new(),
1471        ));
1472        let (message_handler, raw_rx) = channel_message_handler();
1473        let client = WebSocketClient::builder()
1474            .config(WebSocketConfig::builder().url(url).build().unwrap())
1475            .message_handler(message_handler)
1476            .rate_limiter(Arc::clone(&limiter))
1477            .connect()
1478            .await
1479            .unwrap();
1480        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1481        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1482
1483        let mut handler = OKXWsFeedHandler::new(
1484            Arc::new(AtomicBool::new(false)),
1485            cmd_rx,
1486            raw_rx,
1487            out_tx,
1488            AuthTracker::new(),
1489            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1490            get_atomic_clock_realtime(),
1491        );
1492        handler.inner = Some(client);
1493        let (messages_tx, mut messages_rx) = tokio::sync::mpsc::unbounded_channel();
1494
1495        let running = tokio::spawn(async move {
1496            while let Some(message) = handler.next().await {
1497                messages_tx.send(message).unwrap();
1498            }
1499        });
1500
1501        let instrument_id = InstrumentId::from("BTC-USDT.OKX");
1502        let channels = AtomicMap::new();
1503        channels.insert(instrument_id, OKXBookChannel::Book);
1504        let tracker = BookSyncTracker::default();
1505        tracker.record_subscription(instrument_id, time::Instant::now(), SnapshotGate::default());
1506        let recovery = tracker.claim_recovery(instrument_id).unwrap();
1507        assert!(recovery.begin_replacement());
1508
1509        let arg = OKXSubscriptionArg {
1510            channel: OKXWsChannel::Books,
1511            inst_id: Some(Ustr::from("BTC-USDT")),
1512            inst_type: None,
1513            inst_family: None,
1514        };
1515
1516        cmd_tx
1517            .send(HandlerCommand::Subscribe {
1518                args: vec![arg.clone()],
1519            })
1520            .unwrap();
1521
1522        assert_eq!(
1523            tokio::time::timeout(Duration::from_secs(3), wire_rx.recv())
1524                .await
1525                .unwrap()
1526                .unwrap(),
1527            "subscribe"
1528        );
1529
1530        for action in [OKXBookAction::Snapshot, OKXBookAction::Update] {
1531            let message = tokio::time::timeout(Duration::from_secs(3), messages_rx.recv())
1532                .await
1533                .unwrap()
1534                .unwrap();
1535            assert!(
1536                matches!(message, OKXWsMessage::BookData { action: received, .. } if received == action)
1537            );
1538        }
1539
1540        let (completion, result) = tokio::sync::oneshot::channel();
1541        cmd_tx
1542            .send(HandlerCommand::Resubscribe {
1543                subscription: arg.clone(),
1544                cancel: recovery.cancellation.child_token(),
1545                gate: recovery.gate.clone(),
1546                completion,
1547            })
1548            .unwrap();
1549
1550        assert_eq!(
1551            tokio::time::timeout(Duration::from_secs(3), wire_rx.recv())
1552                .await
1553                .unwrap()
1554                .unwrap(),
1555            "unsubscribe"
1556        );
1557
1558        // Paused time keeps the subscribe blocked until the reset completes
1559        tokio::time::pause();
1560        tracker.reset_sequences(&channels, BookChannelScope::Public);
1561
1562        if !snapshot_timeout.is_zero() {
1563            tracker.seed_pending_snapshots(
1564                &channels,
1565                BookChannelScope::Public,
1566                snapshot_timeout,
1567                time::Instant::now(),
1568            );
1569        }
1570
1571        tokio::time::advance(Duration::from_secs(10)).await;
1572        tokio::time::resume();
1573        let _ = tokio::time::timeout(Duration::from_secs(3), result)
1574            .await
1575            .expect("recovery command completes after reconnect");
1576
1577        let resumed = tokio::time::timeout(Duration::from_secs(3), async {
1578            for (action, sequence) in [(OKXBookAction::Snapshot, 100), (OKXBookAction::Update, 101)]
1579            {
1580                let Some(OKXWsMessage::BookData {
1581                    action: received,
1582                    data,
1583                    ..
1584                }) = messages_rx.recv().await
1585                else {
1586                    panic!("expected book data after reconnect");
1587                };
1588
1589                assert_eq!(received, action);
1590                assert_eq!(data[0].seq_id, sequence);
1591                assert_eq!(
1592                    tracker.validate_sequence(
1593                        instrument_id,
1594                        received == OKXBookAction::Snapshot,
1595                        &[(data[0].prev_seq_id, data[0].seq_id)],
1596                        time::Instant::now(),
1597                    ),
1598                    BookSequenceOutcome::Accept,
1599                );
1600            }
1601        })
1602        .await;
1603
1604        // A successful manual subscribe isolates a failure to automatic replay
1605        if resumed.is_err() {
1606            cmd_tx
1607                .send(HandlerCommand::Subscribe { args: vec![arg] })
1608                .unwrap();
1609
1610            for action in [OKXBookAction::Snapshot, OKXBookAction::Update] {
1611                let message = tokio::time::timeout(Duration::from_secs(3), messages_rx.recv())
1612                    .await
1613                    .unwrap()
1614                    .unwrap();
1615                assert!(
1616                    matches!(message, OKXWsMessage::BookData { action: received, .. } if received == action)
1617                );
1618            }
1619        }
1620
1621        cmd_tx.send(HandlerCommand::Disconnect).unwrap();
1622        tokio::time::timeout(Duration::from_secs(3), running)
1623            .await
1624            .unwrap()
1625            .unwrap();
1626        server.abort();
1627
1628        assert!(
1629            resumed.is_ok(),
1630            "reconnect must preserve an in-flight recovery and resume book output"
1631        );
1632    }
1633
1634    #[rstest]
1635    fn test_should_retry_typed_transport_and_timeout_errors() {
1636        assert!(!should_retry_replay_safe_error(&OKXWsError::TransportSend(
1637            SendError::BufferFull
1638        )));
1639        assert!(should_retry_replay_safe_error(&OKXWsError::TransportSend(
1640            SendError::Timeout
1641        )));
1642        assert!(should_retry_replay_safe_error(&OKXWsError::TransportSend(
1643            SendError::ConnectionChanged
1644        )));
1645        assert!(!should_retry_replay_safe_error(&OKXWsError::TransportSend(
1646            SendError::WriteTimeout
1647        )));
1648        assert!(!should_retry_replay_safe_error(&OKXWsError::TransportSend(
1649            SendError::BrokenPipe("connection reset".to_string())
1650        )));
1651        assert!(should_retry_replay_safe_error(
1652            &OKXWsError::OperationTimeout { timeout_ms: 1_000 }
1653        ));
1654        assert!(!should_retry_replay_safe_error(&OKXWsError::NoActiveClient));
1655        assert!(!should_retry_replay_safe_error(
1656            &OKXWsError::HandlerUnavailable("closed".to_string())
1657        ));
1658    }
1659
1660    #[rstest]
1661    fn test_retryability_uses_websocket_error_type_not_message() {
1662        let message = "connection reset".to_string();
1663        let temporary = OKXWsError::OkxError {
1664            error_code: "50011".to_string(),
1665            message: message.clone(),
1666        };
1667        let permanent = OKXWsError::ClientError(message.clone());
1668        let ambiguous = OKXWsError::SendFailed(message);
1669
1670        assert!(should_retry_replay_safe_error(&temporary));
1671        assert!(!should_retry_replay_safe_error(&permanent));
1672        assert!(!should_retry_replay_safe_error(&ambiguous));
1673    }
1674
1675    #[rstest]
1676    fn test_subscription_error_restores_failed_unsubscribe() {
1677        let handler = create_handler();
1678        let arg = OKXWebSocketArg {
1679            channel: OKXWsChannel::Books,
1680            inst_id: Some(Ustr::from("BTC-USD")),
1681            inst_type: None,
1682            inst_family: None,
1683            bar: None,
1684        };
1685        let topic = topic_from_websocket_arg(&arg);
1686        handler.subscriptions_state.mark_subscribe(&topic);
1687        handler.subscriptions_state.confirm_subscribe(&topic);
1688        handler.subscriptions_state.mark_unsubscribe(&topic);
1689
1690        let rejected_subscription =
1691            handler.handle_subscription_error(&arg, "60019", "Unsubscription failed");
1692
1693        assert!(!rejected_subscription);
1694        assert_eq!(handler.subscriptions_state.all_topics(), vec![topic]);
1695        assert!(
1696            handler
1697                .subscriptions_state
1698                .pending_subscribe_topics()
1699                .is_empty()
1700        );
1701        assert!(
1702            handler
1703                .subscriptions_state
1704                .pending_unsubscribe_topics()
1705                .is_empty()
1706        );
1707    }
1708
1709    #[rstest]
1710    fn test_subscription_arg_from_error_message_matches_mainnet_shape() {
1711        let msg = "Wrong URL or channel:books,instId:BTC-USDT-SWAP doesn't exist. Please use the \
1712                   correct URL, channel and parameters referring to API document.";
1713
1714        let arg = subscription_arg_from_error_message(msg).unwrap();
1715
1716        assert_eq!(arg.channel, OKXWsChannel::Books);
1717        assert_eq!(arg.inst_id, Some(Ustr::from("BTC-USDT-SWAP")));
1718        assert_eq!(arg.inst_type, None);
1719        assert_eq!(arg.inst_family, None);
1720        assert_eq!(arg.bar, None);
1721    }
1722
1723    #[rstest]
1724    #[case::book_arg(
1725        Some(OKXWsChannel::Books),
1726        "60018",
1727        "Subscription failed",
1728        log::Level::Warn
1729    )]
1730    #[case::book_from_message(
1731        None,
1732        "60018",
1733        "Wrong URL or channel:books,instId:BTC-USDT-SWAP doesn't exist.",
1734        log::Level::Warn
1735    )]
1736    #[case::retryable(
1737        Some(OKXWsChannel::Tickers),
1738        "50001",
1739        "Service unavailable",
1740        log::Level::Warn
1741    )]
1742    #[case::non_book_arg(
1743        Some(OKXWsChannel::Tickers),
1744        "60018",
1745        "Subscription failed",
1746        log::Level::Error
1747    )]
1748    #[case::non_book_from_message(
1749        None,
1750        "60018",
1751        "Wrong URL or channel:tickers,instId:BTC-USDT-SWAP doesn't exist.",
1752        log::Level::Error
1753    )]
1754    fn test_error_frame_level(
1755        #[case] channel: Option<OKXWsChannel>,
1756        #[case] code: &str,
1757        #[case] msg: &str,
1758        #[case] expected: log::Level,
1759    ) {
1760        let arg = channel.map(|channel| OKXWebSocketArg {
1761            channel,
1762            inst_id: Some(Ustr::from("BTC-USDT-SWAP")),
1763            inst_type: None,
1764            inst_family: None,
1765            bar: None,
1766        });
1767
1768        let level = error_frame_level(arg.as_ref(), code, msg);
1769
1770        assert_eq!(level, expected);
1771    }
1772
1773    #[rstest]
1774    fn test_subscription_error_ignores_non_pending_topic() {
1775        let handler = create_handler();
1776        let arg = OKXWebSocketArg {
1777            channel: OKXWsChannel::Books,
1778            inst_id: Some(Ustr::from("BTC-USDT-SWAP")),
1779            inst_type: None,
1780            inst_family: None,
1781            bar: None,
1782        };
1783
1784        let rejected_subscription =
1785            handler.handle_subscription_error(&arg, "60018", "Subscription failed");
1786
1787        assert!(!rejected_subscription);
1788        assert!(handler.subscriptions_state.all_topics().is_empty());
1789        assert!(
1790            handler
1791                .subscriptions_state
1792                .pending_subscribe_topics()
1793                .is_empty()
1794        );
1795        assert!(
1796            handler
1797                .subscriptions_state
1798                .pending_unsubscribe_topics()
1799                .is_empty()
1800        );
1801    }
1802
1803    #[derive(serde::Deserialize, Debug, PartialEq)]
1804    struct ParseArrayItem {
1805        value: i64,
1806    }
1807
1808    #[rstest]
1809    fn test_parse_array_items_keeps_good_items_when_one_fails() {
1810        let data = json!([
1811            {"value": 1},
1812            {"value": "not a number"},
1813            {"value": 3},
1814        ]);
1815
1816        let parsed: Vec<ParseArrayItem> =
1817            parse_array_items(data, "test", false).expect("non-empty");
1818        assert_eq!(
1819            parsed,
1820            vec![ParseArrayItem { value: 1 }, ParseArrayItem { value: 3 }],
1821        );
1822    }
1823
1824    #[rstest]
1825    fn test_parse_array_items_returns_none_when_payload_not_array() {
1826        let data = json!({"not": "an array"});
1827        let parsed: Option<Vec<ParseArrayItem>> = parse_array_items(data, "test", false);
1828        assert!(parsed.is_none());
1829    }
1830
1831    #[rstest]
1832    fn test_parse_array_items_returns_none_when_all_items_fail() {
1833        let data = json!([{"value": "bad"}]);
1834        let parsed: Option<Vec<ParseArrayItem>> = parse_array_items(data, "test", false);
1835        assert!(parsed.is_none());
1836    }
1837
1838    #[rstest]
1839    fn test_route_instruments_keeps_valid_items_when_one_item_fails() {
1840        let handler = create_handler();
1841        let mut frame: Value =
1842            serde_json::from_str(&load_test_json("ws_instruments.json")).expect("valid fixture");
1843        let data = frame
1844            .get_mut("data")
1845            .and_then(Value::as_array_mut)
1846            .expect("data array");
1847        let mut invalid_item = data[0].clone();
1848        invalid_item["tickSz"] = json!(7);
1849        data.insert(0, invalid_item);
1850
1851        let arg: OKXWebSocketArg = serde_json::from_value(frame["arg"].clone()).expect("valid arg");
1852        let msg = handler
1853            .route_data_message(arg, frame["data"].clone())
1854            .expect("instruments message");
1855
1856        match msg {
1857            OKXWsMessage::Instruments(instruments) => {
1858                assert_eq!(instruments.len(), 1);
1859                assert_eq!(instruments[0].inst_id, "BTC-USDT-SWAP");
1860            }
1861            other => panic!("Expected Instruments, was {other:?}"),
1862        }
1863    }
1864
1865    #[rstest]
1866    fn test_route_instruments_prefers_rpi_over_legacy_alias() {
1867        let handler = create_handler();
1868        let mut frame: Value =
1869            serde_json::from_str(&load_test_json("ws_instruments.json")).expect("valid fixture");
1870        let instrument = &mut frame["data"][0];
1871        instrument["rpi"] = json!("2");
1872        instrument["elp"] = json!("1");
1873
1874        let arg: OKXWebSocketArg = serde_json::from_value(frame["arg"].clone()).expect("valid arg");
1875        let msg = handler
1876            .route_data_message(arg, frame["data"].clone())
1877            .expect("instruments message");
1878
1879        match msg {
1880            OKXWsMessage::Instruments(instruments) => {
1881                assert_eq!(instruments.len(), 1);
1882                assert_eq!(instruments[0].rpi, Some(OKXRpiPermission::Permitted));
1883            }
1884            other => panic!("Expected Instruments, was {other:?}"),
1885        }
1886    }
1887
1888    #[rstest]
1889    fn test_route_liquidation_warnings() {
1890        let handler = create_handler();
1891        let frame: Value = serde_json::from_str(&load_test_json("ws_liquidation_warning.json"))
1892            .expect("valid fixture");
1893
1894        let arg: OKXWebSocketArg = serde_json::from_value(frame["arg"].clone()).expect("valid arg");
1895        let msg = handler
1896            .route_data_message(arg, frame["data"].clone())
1897            .expect("liquidation warning message");
1898
1899        match msg {
1900            OKXWsMessage::LiquidationWarnings(warnings) => {
1901                assert_eq!(warnings.len(), 1);
1902                assert_eq!(warnings[0].inst_id, "BTC-USDT-SWAP");
1903                assert_eq!(warnings[0].mgn_ratio, "0.62");
1904            }
1905            other => panic!("Expected LiquidationWarnings, was {other:?}"),
1906        }
1907    }
1908
1909    #[tokio::test]
1910    async fn bound_send_rejects_reconnecting_epoch_before_waiting() {
1911        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
1912        let url = format!("ws://{}/", listener.local_addr().unwrap());
1913
1914        let server = tokio::spawn(async move {
1915            let (socket, _) = listener.accept().await.unwrap();
1916            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
1917            let mut frames = Vec::new();
1918            while let Some(Ok(Message::Text(text))) = socket.next().await {
1919                frames.push(text.to_string());
1920            }
1921
1922            frames
1923        });
1924
1925        let (message_handler, raw_rx) = channel_message_handler();
1926        let client = WebSocketClient::builder()
1927            .config(WebSocketConfig::builder().url(url).build().unwrap())
1928            .message_handler(message_handler)
1929            .connect()
1930            .await
1931            .unwrap();
1932        let mode = client.connection_mode_atomic();
1933        let epoch = client.connection_epoch();
1934        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1935        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
1936
1937        let mut handler = OKXWsFeedHandler::new(
1938            Arc::new(AtomicBool::new(false)),
1939            cmd_rx,
1940            raw_rx,
1941            out_tx,
1942            AuthTracker::new(),
1943            SubscriptionState::new(OKX_WS_TOPIC_DELIMITER),
1944            get_atomic_clock_realtime(),
1945        );
1946        handler.inner = Some(client);
1947        handler
1948            .send_on_connection("active-control".to_string(), None)
1949            .await
1950            .unwrap();
1951
1952        tokio::time::pause();
1953        mode.store(ConnectionMode::Reconnect.as_u8(), Ordering::SeqCst);
1954
1955        let (ready, result) = {
1956            let mut send =
1957                Box::pin(handler.send_on_connection("cancel-in-reconnect".to_string(), None));
1958
1959            match futures_util::poll!(send.as_mut()) {
1960                std::task::Poll::Ready(result) => (true, result),
1961                std::task::Poll::Pending => {
1962                    mode.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
1963                    tokio::time::advance(Duration::from_millis(100)).await;
1964                    (false, send.await)
1965                }
1966            }
1967        };
1968
1969        mode.store(ConnectionMode::Active.as_u8(), Ordering::SeqCst);
1970        tokio::time::resume();
1971        let final_epoch = handler.inner.as_ref().unwrap().connection_epoch();
1972        handler.inner.as_ref().unwrap().disconnect().await;
1973        let frames = tokio::time::timeout(Duration::from_secs(2), server)
1974            .await
1975            .expect("peer did not close")
1976            .unwrap();
1977
1978        assert_eq!(frames, vec!["active-control"]);
1979        assert_eq!(final_epoch, epoch);
1980        assert!(
1981            ready,
1982            "reconnecting send must fail without waiting for recovery"
1983        );
1984        assert!(
1985            matches!(
1986                result,
1987                Err(OKXWsError::TransportSend(SendError::ConnectionChanged))
1988            ),
1989            "expected rejection for a reconnecting epoch, received {result:?}"
1990        );
1991    }
1992}