nautilus_derive/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//! Inner I/O feed handler for the Derive WebSocket transport.
17//!
18//! The handler owns the [`WebSocketClient`] exclusively and runs in a dedicated
19//! Tokio task. The outer [`super::client::DeriveWebSocketClient`] talks to it
20//! via a command channel and consumes a stream of [`DeriveWsMessage`] events.
21//!
22//! Each outbound JSON-RPC request is registered in a `pending` map keyed by the
23//! correlator `id`. When the venue echoes the id on a response frame, the
24//! matching oneshot is fulfilled with `result` or the JSON-RPC error.
25
26use std::{
27    sync::{
28        Arc,
29        atomic::{AtomicBool, AtomicU64, Ordering},
30    },
31    time::Duration,
32};
33
34use ahash::AHashMap;
35use nautilus_core::string::secret::SecretString;
36use nautilus_live::task::{SharedTaskSlot, TaskJoinOutcome};
37use nautilus_network::{
38    RECONNECTED,
39    websocket::{AuthTracker, WebSocketClient},
40};
41use serde_json::Value;
42use tokio_tungstenite::tungstenite::Message;
43
44use super::{
45    client::UNAUTHENTICATED_CONNECTION_EPOCH,
46    error::DeriveWsError,
47    messages::{DeriveWsChannel, DeriveWsFrame, WsSubscribeParams, WsSubscriptionPayload},
48};
49use crate::http::models::JsonRpcRequest;
50
51/// Outbound commands the outer client sends to the inner handler.
52#[derive(Debug)]
53pub(super) enum HandlerCommand {
54    /// Hand the active [`WebSocketClient`] to the handler.
55    SetClient(WebSocketClient),
56    /// Send a JSON-RPC request and resolve the oneshot when the venue replies.
57    /// `params` is pre-serialized JSON so the handler stays agnostic to the
58    /// per-method param types (login, subscribe, signed `private/*` bodies).
59    Request {
60        method: &'static str,
61        params: SecretString,
62        connection_epoch: Option<u64>,
63        response_tx: tokio::sync::oneshot::Sender<Result<Value, DeriveWsError>>,
64    },
65    /// Gracefully tear down the WebSocket connection.
66    Disconnect,
67}
68
69/// Events emitted by the handler for the outer client and downstream consumers.
70#[derive(Debug, Clone)]
71pub enum DeriveWsMessage {
72    /// `public/login` succeeded. Consumed by the client's spawn loop to drive
73    /// resubscription; not forwarded to data/execution layers.
74    Authenticated,
75    /// Underlying transport reconnected; outer client triggers re-login and
76    /// resubscribes the tracked channels.
77    Reconnected,
78    /// Re-login or subscription replay exhausted its retry budget.
79    SessionRecoveryFailed(String),
80    /// Channel update pushed by the venue.
81    Subscription(WsSubscriptionPayload),
82}
83
84#[derive(Debug)]
85struct SendCommand {
86    id: u64,
87    token: u64,
88    connection_epoch: u64,
89    payload: SecretString,
90}
91
92#[derive(Debug)]
93struct SendFailure {
94    id: u64,
95    token: u64,
96    reason: String,
97}
98
99#[derive(Debug)]
100struct PendingRequest {
101    token: u64,
102    response_tx: tokio::sync::oneshot::Sender<Result<Value, DeriveWsError>>,
103}
104
105/// Inner I/O loop. Lives in a Tokio task spawned by
106/// [`super::client::DeriveWebSocketClient::connect`].
107pub(super) struct FeedHandler {
108    signal: Arc<AtomicBool>,
109    client: Option<Arc<WebSocketClient>>,
110    cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
111    cmd_closed: bool,
112    raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
113    raw_closed: bool,
114    next_id: Arc<AtomicU64>,
115    next_send_token: u64,
116    pending: AHashMap<u64, PendingRequest>,
117    send_tx: Option<tokio::sync::mpsc::UnboundedSender<SendCommand>>,
118    send_failure_rx: tokio::sync::mpsc::UnboundedReceiver<SendFailure>,
119    send_task: Arc<SharedTaskSlot<()>>,
120    auth_tracker: AuthTracker,
121    authenticated_epoch: Arc<AtomicU64>,
122}
123
124impl FeedHandler {
125    pub(super) fn new_with_send_task(
126        signal: Arc<AtomicBool>,
127        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
128        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
129        next_id: Arc<AtomicU64>,
130        auth_tracker: AuthTracker,
131        authenticated_epoch: Arc<AtomicU64>,
132        send_task: Arc<SharedTaskSlot<()>>,
133    ) -> Self {
134        let (_, send_failure_rx) = tokio::sync::mpsc::unbounded_channel();
135        Self {
136            signal,
137            client: None,
138            cmd_rx,
139            cmd_closed: false,
140            raw_rx,
141            raw_closed: false,
142            next_id,
143            next_send_token: 0,
144            pending: AHashMap::new(),
145            send_tx: None,
146            send_failure_rx,
147            send_task,
148            auth_tracker,
149            authenticated_epoch,
150        }
151    }
152
153    /// Drains the next event from the underlying channels, processes it, and
154    /// returns the resulting outbound message (if any). Returns `None` when
155    /// the handler is shutting down or both channels closed.
156    pub(super) async fn next(&mut self) -> Option<DeriveWsMessage> {
157        loop {
158            tokio::select! {
159                cmd = self.cmd_rx.recv(), if !self.cmd_closed => {
160                    match cmd {
161                        None => {
162                            self.cmd_closed = true;
163
164                            if self.raw_closed {
165                                self.shutdown_send_path(
166                                    "WebSocket handler stopped before response was received",
167                                )
168                                .await;
169                                return None;
170                            }
171                        }
172                        Some(HandlerCommand::SetClient(client)) => {
173                            log::debug!("Setting WebSocket client in Derive handler");
174                            let client = Arc::new(client);
175                            self.stop_send_worker().await;
176                            self.start_send_worker(Arc::clone(&client));
177                            self.client = Some(client);
178                        }
179                        Some(HandlerCommand::Request {
180                            method,
181                            params,
182                            connection_epoch,
183                            response_tx,
184                        }) => {
185                            self.dispatch_request(method, params, connection_epoch, response_tx);
186                        }
187                        Some(HandlerCommand::Disconnect) => {
188                            log::debug!("Derive handler received disconnect command");
189                            self.shutdown_send_path(
190                                "WebSocket disconnected before response was received",
191                            )
192                            .await;
193
194                            if let Some(ref client) = self.client {
195                                client.disconnect().await;
196                            }
197                            self.signal.store(true, Ordering::SeqCst);
198                            return None;
199                        }
200                    }
201                }
202
203                raw = self.raw_rx.recv(), if !self.raw_closed => {
204                    match raw {
205                        None => {
206                            self.raw_closed = true;
207
208                            if self.cmd_closed {
209                                self.shutdown_send_path(
210                                    "WebSocket handler stopped before response was received",
211                                )
212                                .await;
213                                return None;
214                            }
215                        }
216                        Some(Message::Text(text)) => {
217                            if text.as_str() == RECONNECTED {
218                                log::info!("Derive WebSocket reconnected sentinel received");
219                                self.auth_tracker.invalidate();
220                                self.authenticated_epoch.store(
221                                    UNAUTHENTICATED_CONNECTION_EPOCH,
222                                    Ordering::Release,
223                                );
224                                self.restart_send_worker(
225                                    "WebSocket reconnected before response was received",
226                                )
227                                .await;
228                                return Some(DeriveWsMessage::Reconnected);
229                            }
230
231                            match DeriveWsFrame::parse(&text) {
232                                Ok(DeriveWsFrame::Response { id, result, error }) => {
233                                    if let Some(pending) = self.pending.remove(&id) {
234                                        let outcome = match (result, error) {
235                                            (_, Some(err)) => Err(DeriveWsError::JsonRpc {
236                                                code: err.code,
237                                                message: err.message,
238                                                data: err.data,
239                                            }),
240                                            (Some(value), None) => Ok(value),
241                                            (None, None) => Ok(Value::Null),
242                                        };
243                                        let _ = pending.response_tx.send(outcome);
244                                    } else {
245                                        log::debug!(
246                                            "Derive WebSocket response with unknown id={id} dropped",
247                                        );
248                                    }
249                                }
250                                Ok(DeriveWsFrame::Subscription(payload)) => {
251                                    return Some(DeriveWsMessage::Subscription(payload));
252                                }
253                                Ok(DeriveWsFrame::UncorrelatedError(error)) => {
254                                    self.fail_uncorrelated_error(error);
255                                }
256                                Ok(DeriveWsFrame::Unknown(value)) => {
257                                    log::debug!("Derive WebSocket unknown frame: {value}");
258                                }
259                                Err(e) => {
260                                    log::error!(
261                                        "Derive WebSocket frame parse error: {e}, text: {text}",
262                                    );
263                                }
264                            }
265                        }
266                        Some(Message::Ping(data)) => {
267                            if let Some(ref client) = self.client
268                                && let Err(e) = client.send_pong(data.to_vec()).await {
269                                log::error!("Derive WebSocket send_pong failed: {e}");
270                            }
271                        }
272                        Some(Message::Close(_)) => {
273                            log::debug!("Derive WebSocket close frame received");
274                            self.shutdown_send_path(
275                                "WebSocket closed before response was received",
276                            )
277                            .await;
278                            return None;
279                        }
280                        Some(_) => {}
281                    }
282                }
283
284                Some(failure) = self.send_failure_rx.recv() => {
285                    self.handle_send_failure(failure);
286                }
287
288                else => {
289                    log::debug!("Derive handler shutting down: channels closed");
290                    self.shutdown_send_path(
291                        "WebSocket handler stopped before response was received",
292                    )
293                    .await;
294                    return None;
295                }
296            }
297        }
298    }
299
300    fn dispatch_request(
301        &mut self,
302        method: &'static str,
303        params: SecretString,
304        connection_epoch: Option<u64>,
305        response_tx: tokio::sync::oneshot::Sender<Result<Value, DeriveWsError>>,
306    ) {
307        let Some(client) = self.client.as_ref() else {
308            let _ = response_tx.send(Err(DeriveWsError::NotConnected));
309            return;
310        };
311        let current_epoch = client.connection_epoch();
312        let connection_epoch = if method.starts_with("private/") {
313            let authenticated_epoch = self.authenticated_epoch.load(Ordering::Acquire);
314            if !client.connection_mode().is_active() || authenticated_epoch != current_epoch {
315                let _ = response_tx.send(Err(DeriveWsError::Authentication {
316                    operation: method.to_string(),
317                    reason: "WebSocket session is not authenticated".to_string(),
318                }));
319                return;
320            }
321            authenticated_epoch
322        } else {
323            connection_epoch.unwrap_or(current_epoch)
324        };
325
326        if !client.connection_mode().is_active() || connection_epoch != current_epoch {
327            let _ = response_tx.send(Err(DeriveWsError::transport(format!(
328                "connection changed before `{method}` was sent",
329            ))));
330            return;
331        }
332
333        self.enqueue_request(method, params, connection_epoch, response_tx);
334    }
335
336    fn enqueue_request(
337        &mut self,
338        method: &'static str,
339        params_json: SecretString,
340        connection_epoch: u64,
341        response_tx: tokio::sync::oneshot::Sender<Result<Value, DeriveWsError>>,
342    ) {
343        let Some(send_tx) = self.send_tx.clone() else {
344            let _ = response_tx.send(Err(DeriveWsError::NotConnected));
345            return;
346        };
347        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
348        let token = self.next_send_token;
349        self.next_send_token = self.next_send_token.wrapping_add(1);
350        let params: Value = match serde_json::from_str(params_json.expose_secret()) {
351            Ok(params) => params,
352            Err(e) => {
353                let _ = response_tx.send(Err(DeriveWsError::Serde(e)));
354                return;
355            }
356        };
357        drop(params_json);
358        let request = JsonRpcRequest::new(id, method, params);
359        let payload = match serde_json::to_string(&request) {
360            Ok(payload) => SecretString::from(payload),
361            Err(e) => {
362                let _ = response_tx.send(Err(DeriveWsError::Serde(e)));
363                return;
364            }
365        };
366        self.pending
367            .insert(id, PendingRequest { token, response_tx });
368        log::debug!("Derive WebSocket sending `{method}` id={id}");
369        if let Err(e) = send_tx.send(SendCommand {
370            id,
371            token,
372            connection_epoch,
373            payload,
374        }) && self
375            .pending
376            .get(&id)
377            .is_some_and(|pending| pending.token == token)
378            && let Some(pending) = self.pending.remove(&id)
379        {
380            let _ = pending
381                .response_tx
382                .send(Err(DeriveWsError::transport(format!(
383                    "failed to queue WebSocket request: {e}",
384                ))));
385        }
386    }
387
388    fn start_send_worker(&mut self, client: Arc<WebSocketClient>) {
389        if !self.send_task.is_empty() {
390            log::error!("Cannot start Derive WebSocket send worker while the prior task is owned");
391            return;
392        }
393
394        let (send_tx, send_rx) = tokio::sync::mpsc::unbounded_channel();
395        let (failure_tx, failure_rx) = tokio::sync::mpsc::unbounded_channel();
396        self.send_tx = Some(send_tx);
397        self.send_failure_rx = failure_rx;
398
399        if let Err(e) = self
400            .send_task
401            .spawn(run_send_worker(client, send_rx, failure_tx))
402        {
403            self.send_tx.take();
404            log::error!("Failed to start Derive WebSocket send worker: {e}");
405        }
406    }
407
408    async fn stop_send_worker(&mut self) {
409        self.send_tx.take();
410
411        if let Some(outcome) = self
412            .send_task
413            .finish(Duration::from_secs(1), Duration::from_secs(2))
414            .await
415        {
416            match outcome {
417                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
418                TaskJoinOutcome::Failed(e) => {
419                    log::error!("Derive WebSocket send worker failed: {e}");
420                }
421                TaskJoinOutcome::Incomplete => {
422                    log::error!("Derive WebSocket send worker did not stop after abort");
423                }
424            }
425        }
426    }
427
428    async fn restart_send_worker(&mut self, reason: &str) {
429        self.stop_send_worker().await;
430        self.fail_pending(reason);
431
432        if let Some(client) = self.client.clone() {
433            self.start_send_worker(client);
434        }
435    }
436
437    async fn shutdown_send_path(&mut self, reason: &str) {
438        self.stop_send_worker().await;
439        self.fail_pending(reason);
440    }
441
442    fn handle_send_failure(&mut self, failure: SendFailure) {
443        if !self
444            .pending
445            .get(&failure.id)
446            .is_some_and(|pending| pending.token == failure.token)
447        {
448            return;
449        }
450
451        if let Some(pending) = self.pending.remove(&failure.id) {
452            let _ = pending
453                .response_tx
454                .send(Err(DeriveWsError::transport(failure.reason)));
455        }
456    }
457
458    fn fail_uncorrelated_error(&mut self, error: crate::http::models::JsonRpcError) {
459        if self.pending.len() != 1 {
460            log::warn!(
461                "Derive WebSocket uncorrelated JSON-RPC error with {} pending requests: code={}, message={}",
462                self.pending.len(),
463                error.code,
464                error.message,
465            );
466            return;
467        }
468
469        if let Some((_, pending)) = self.pending.drain().next() {
470            let _ = pending.response_tx.send(Err(DeriveWsError::JsonRpc {
471                code: error.code,
472                message: error.message,
473                data: error.data,
474            }));
475        }
476    }
477
478    fn fail_pending(&mut self, reason: &str) {
479        if self.pending.is_empty() {
480            return;
481        }
482        log::debug!(
483            "Failing {} pending Derive WebSocket request(s): {reason}",
484            self.pending.len(),
485        );
486
487        for (_, pending) in self.pending.drain() {
488            let _ = pending
489                .response_tx
490                .send(Err(DeriveWsError::transport(reason.to_string())));
491        }
492    }
493}
494
495impl Drop for FeedHandler {
496    fn drop(&mut self) {
497        self.send_tx.take();
498        self.send_task.abort();
499    }
500}
501
502async fn run_send_worker(
503    client: Arc<WebSocketClient>,
504    mut send_rx: tokio::sync::mpsc::UnboundedReceiver<SendCommand>,
505    failure_tx: tokio::sync::mpsc::UnboundedSender<SendFailure>,
506) {
507    while let Some(command) = send_rx.recv().await {
508        if let Err(e) = client
509            .send_text_on_connection(
510                command.payload.expose_secret().to_owned(),
511                None,
512                command.connection_epoch,
513            )
514            .await
515            && failure_tx
516                .send(SendFailure {
517                    id: command.id,
518                    token: command.token,
519                    reason: e.to_string(),
520                })
521                .is_err()
522        {
523            return;
524        }
525    }
526}
527
528/// Builds `subscribe` params from a single channel topic.
529#[must_use]
530pub(super) fn subscribe_params(channel: DeriveWsChannel) -> WsSubscribeParams {
531    WsSubscribeParams {
532        channels: vec![channel],
533    }
534}
535
536/// Convenience wrapper that produces the `subscribe` params for the
537/// `ticker_slim.{instrument_name}.{interval}` channel.
538#[must_use]
539pub(super) fn ticker_subscribe_params(instrument_name: &str, interval: &str) -> WsSubscribeParams {
540    subscribe_params(DeriveWsChannel::ticker_slim(instrument_name, interval))
541}
542
543/// Convenience wrapper that produces the `subscribe` params for the
544/// `orderbook.{instrument_name}.{group}.{depth}` channel.
545#[must_use]
546pub(super) fn orderbook_subscribe_params(
547    instrument_name: &str,
548    group: &str,
549    depth: &str,
550) -> WsSubscribeParams {
551    subscribe_params(DeriveWsChannel::orderbook(instrument_name, group, depth))
552}
553
554/// Convenience wrapper that produces the `subscribe` params for the
555/// `trades.{instrument_type}.{currency}` channel.
556#[must_use]
557pub(super) fn trades_subscribe_params(instrument_type: &str, currency: &str) -> WsSubscribeParams {
558    subscribe_params(DeriveWsChannel::trades(instrument_type, currency))
559}
560
561#[cfg(test)]
562mod tests {
563    use nautilus_core::string::secret::zeroize_json_value;
564    use rstest::rstest;
565    use serde_json::json;
566
567    use super::*;
568
569    fn unauthenticated_epoch() -> Arc<AtomicU64> {
570        Arc::new(AtomicU64::new(u64::MAX))
571    }
572
573    fn secret_json(mut value: Value) -> SecretString {
574        let serialized = SecretString::from(value.to_string());
575        zeroize_json_value(&mut value);
576        serialized
577    }
578
579    #[rstest]
580    fn test_queued_request_debug_redacts_payload() {
581        let secret = "signature-secret";
582        let (response_tx, _response_rx) = tokio::sync::oneshot::channel();
583        let command = HandlerCommand::Request {
584            method: "public/login",
585            params: SecretString::from(format!(r#"{{"signature":"{secret}"}}"#)),
586            connection_epoch: None,
587            response_tx,
588        };
589        let send = SendCommand {
590            id: 1,
591            token: 1,
592            connection_epoch: 1,
593            payload: SecretString::from(format!(r#"{{"signature":"{secret}"}}"#)),
594        };
595
596        let debug = format!("{command:?} {send:?}");
597
598        assert!(!debug.contains(secret));
599    }
600
601    fn feed_handler(
602        signal: Arc<AtomicBool>,
603        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
604        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
605        next_id: Arc<AtomicU64>,
606        auth_tracker: AuthTracker,
607        authenticated_epoch: Arc<AtomicU64>,
608    ) -> FeedHandler {
609        FeedHandler::new_with_send_task(
610            signal,
611            cmd_rx,
612            raw_rx,
613            next_id,
614            auth_tracker,
615            authenticated_epoch,
616            Arc::new(SharedTaskSlot::new()),
617        )
618    }
619
620    #[rstest]
621    fn test_subscribe_params_carries_single_channel() {
622        let params = subscribe_params(DeriveWsChannel::ticker_slim("ETH-PERP", "1000"));
623        assert_eq!(
624            params.channels,
625            vec![DeriveWsChannel::ticker_slim("ETH-PERP", "1000")],
626        );
627    }
628
629    #[rstest]
630    fn test_ticker_subscribe_params_formats_topic() {
631        let params = ticker_subscribe_params("ETH-PERP", "1000");
632        assert_eq!(
633            params.channels,
634            vec![DeriveWsChannel::ticker_slim("ETH-PERP", "1000")],
635        );
636    }
637
638    #[rstest]
639    fn test_orderbook_subscribe_params_formats_topic() {
640        let params = orderbook_subscribe_params("ETH-PERP", "1", "10");
641        assert_eq!(
642            params.channels,
643            vec![DeriveWsChannel::orderbook("ETH-PERP", "1", "10")],
644        );
645    }
646
647    #[rstest]
648    fn test_trades_subscribe_params_formats_topic() {
649        let params = trades_subscribe_params("perp", "ETH");
650        assert_eq!(
651            params.channels,
652            vec![DeriveWsChannel::trades("perp", "ETH")],
653        );
654    }
655
656    #[rstest]
657    #[tokio::test]
658    async fn test_dispatch_request_without_client_returns_not_connected() {
659        // Requests issued before SetClient must fail fast rather than hang.
660        let signal = Arc::new(AtomicBool::new(false));
661        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
662        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
663        let next_id = Arc::new(AtomicU64::new(1));
664        let auth_tracker = AuthTracker::new();
665        let mut handler = feed_handler(
666            signal,
667            cmd_rx,
668            raw_rx,
669            next_id,
670            auth_tracker,
671            unauthenticated_epoch(),
672        );
673
674        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
675        let params =
676            secret_json(serde_json::to_value(WsSubscribeParams { channels: vec![] }).unwrap());
677        handler.dispatch_request("public/login", params, None, response_tx);
678
679        let outcome = response_rx.await.expect("oneshot resolved");
680        match outcome {
681            Err(DeriveWsError::NotConnected) => {}
682            other => panic!("expected NotConnected, was {other:?}"),
683        }
684    }
685
686    #[rstest]
687    #[tokio::test]
688    async fn test_dispatch_registers_pending_requests_before_ordered_queueing() {
689        let signal = Arc::new(AtomicBool::new(false));
690        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
691        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
692        let next_id = Arc::new(AtomicU64::new(1));
693        let auth_tracker = AuthTracker::new();
694        let mut handler = feed_handler(
695            signal,
696            cmd_rx,
697            raw_rx,
698            next_id,
699            auth_tracker,
700            unauthenticated_epoch(),
701        );
702        let (send_tx, mut send_rx) = tokio::sync::mpsc::unbounded_channel();
703        handler.send_tx = Some(send_tx);
704
705        let (first_tx, _first_rx) = tokio::sync::oneshot::channel();
706        let (second_tx, _second_rx) = tokio::sync::oneshot::channel();
707        handler.enqueue_request("first", secret_json(json!({"sequence": 1})), 0, first_tx);
708        handler.enqueue_request("second", secret_json(json!({"sequence": 2})), 0, second_tx);
709
710        let first = send_rx.recv().await.expect("first queued send");
711        let second = send_rx.recv().await.expect("second queued send");
712        let first_payload: Value = serde_json::from_str(first.payload.expose_secret()).unwrap();
713        let second_payload: Value = serde_json::from_str(second.payload.expose_secret()).unwrap();
714
715        assert_eq!(first.id, 1);
716        assert_eq!(second.id, 2);
717        assert_eq!(first_payload["method"], "first");
718        assert_eq!(second_payload["method"], "second");
719        assert_eq!(handler.pending.get(&first.id).unwrap().token, first.token);
720        assert_eq!(handler.pending.get(&second.id).unwrap().token, second.token);
721    }
722
723    #[rstest]
724    #[tokio::test]
725    async fn test_late_send_failure_does_not_remove_reused_request_id() {
726        let signal = Arc::new(AtomicBool::new(false));
727        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
728        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
729        let next_id = Arc::new(AtomicU64::new(1));
730        let auth_tracker = AuthTracker::new();
731        let mut handler = feed_handler(
732            signal,
733            cmd_rx,
734            raw_rx,
735            Arc::clone(&next_id),
736            auth_tracker,
737            unauthenticated_epoch(),
738        );
739        let (send_tx, mut send_rx) = tokio::sync::mpsc::unbounded_channel();
740        handler.send_tx = Some(send_tx);
741
742        let (old_tx, old_rx) = tokio::sync::oneshot::channel();
743        handler.enqueue_request("first", secret_json(json!({})), 0, old_tx);
744        let old_send = send_rx.recv().await.expect("old queued send");
745        let old_pending = handler.pending.remove(&old_send.id).unwrap();
746        old_pending.response_tx.send(Ok(Value::Null)).unwrap();
747        old_rx.await.unwrap().unwrap();
748
749        next_id.store(old_send.id, Ordering::Relaxed);
750        let (new_tx, new_rx) = tokio::sync::oneshot::channel();
751        handler.enqueue_request("second", secret_json(json!({})), 0, new_tx);
752        let new_send = send_rx.recv().await.expect("new queued send");
753        handler.handle_send_failure(SendFailure {
754            id: old_send.id,
755            token: old_send.token,
756            reason: "late failure".to_string(),
757        });
758
759        assert_eq!(old_send.id, new_send.id);
760        assert_ne!(old_send.token, new_send.token);
761        assert_eq!(
762            handler.pending.get(&new_send.id).unwrap().token,
763            new_send.token,
764        );
765
766        handler.handle_send_failure(SendFailure {
767            id: new_send.id,
768            token: new_send.token,
769            reason: "current failure".to_string(),
770        });
771        let error = new_rx.await.unwrap().expect_err("current request failed");
772        assert!(error.to_string().contains("current failure"));
773    }
774
775    #[rstest]
776    #[tokio::test]
777    async fn test_shutdown_aborts_send_worker_and_drains_pending_requests() {
778        let signal = Arc::new(AtomicBool::new(false));
779        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
780        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
781        let next_id = Arc::new(AtomicU64::new(1));
782        let auth_tracker = AuthTracker::new();
783        let mut handler = feed_handler(
784            signal,
785            cmd_rx,
786            raw_rx,
787            next_id,
788            auth_tracker,
789            unauthenticated_epoch(),
790        );
791        let (send_tx, _send_rx) = tokio::sync::mpsc::unbounded_channel();
792        handler.send_tx = Some(send_tx);
793
794        let task = tokio::spawn(std::future::pending::<()>());
795        let abort_handle = task.abort_handle();
796        handler.send_task.insert(task);
797
798        let (first_tx, first_rx) = tokio::sync::oneshot::channel();
799        let (second_tx, second_rx) = tokio::sync::oneshot::channel();
800        handler.enqueue_request("first", secret_json(json!({})), 0, first_tx);
801        handler.enqueue_request("second", secret_json(json!({})), 0, second_tx);
802        handler.shutdown_send_path("disconnect requested").await;
803        tokio::task::yield_now().await;
804
805        let first_error = first_rx.await.unwrap().expect_err("first request failed");
806        let second_error = second_rx.await.unwrap().expect_err("second request failed");
807        assert!(abort_handle.is_finished());
808        assert!(handler.send_tx.is_none());
809        assert!(handler.send_task.is_empty());
810        assert!(handler.pending.is_empty());
811        assert!(first_error.to_string().contains("disconnect requested"));
812        assert!(second_error.to_string().contains("disconnect requested"));
813    }
814
815    #[rstest]
816    #[tokio::test]
817    async fn test_next_stops_when_input_channels_close_with_failure_channel_open() {
818        let signal = Arc::new(AtomicBool::new(false));
819        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
820        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
821        let next_id = Arc::new(AtomicU64::new(1));
822        let auth_tracker = AuthTracker::new();
823        let mut handler = feed_handler(
824            signal,
825            cmd_rx,
826            raw_rx,
827            next_id,
828            auth_tracker,
829            unauthenticated_epoch(),
830        );
831        let (_failure_tx, failure_rx) = tokio::sync::mpsc::unbounded_channel();
832        handler.send_failure_rx = failure_rx;
833        drop(cmd_tx);
834        drop(raw_tx);
835
836        let outcome = tokio::time::timeout(std::time::Duration::from_millis(100), handler.next())
837            .await
838            .expect("handler stopped after both input channels closed");
839
840        assert!(outcome.is_none());
841    }
842
843    #[rstest]
844    #[tokio::test]
845    async fn test_uncorrelated_error_fails_only_pending_request() {
846        let signal = Arc::new(AtomicBool::new(false));
847        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
848        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
849        let next_id = Arc::new(AtomicU64::new(1));
850        let auth_tracker = AuthTracker::new();
851        let mut handler = feed_handler(
852            signal,
853            cmd_rx,
854            raw_rx,
855            next_id,
856            auth_tracker,
857            unauthenticated_epoch(),
858        );
859        let (send_tx, _send_rx) = tokio::sync::mpsc::unbounded_channel();
860        handler.send_tx = Some(send_tx);
861        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
862        handler.enqueue_request("request", secret_json(json!({})), 0, response_tx);
863
864        handler.fail_uncorrelated_error(crate::http::models::JsonRpcError {
865            code: -32700,
866            message: "Parse error".to_string(),
867            data: Some(json!("invalid JSON")),
868        });
869        let error = response_rx.await.unwrap().expect_err("request failed");
870
871        match error {
872            DeriveWsError::JsonRpc {
873                code,
874                message,
875                data,
876            } => {
877                assert_eq!(code, -32700);
878                assert_eq!(message, "Parse error");
879                assert_eq!(data, Some(json!("invalid JSON")));
880            }
881            other => panic!("expected JsonRpc, was {other:?}"),
882        }
883        assert!(handler.pending.is_empty());
884    }
885
886    #[rstest]
887    #[tokio::test]
888    async fn test_uncorrelated_error_does_not_guess_between_pending_requests() {
889        let signal = Arc::new(AtomicBool::new(false));
890        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
891        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
892        let next_id = Arc::new(AtomicU64::new(1));
893        let auth_tracker = AuthTracker::new();
894        let mut handler = feed_handler(
895            signal,
896            cmd_rx,
897            raw_rx,
898            next_id,
899            auth_tracker,
900            unauthenticated_epoch(),
901        );
902        let (send_tx, _send_rx) = tokio::sync::mpsc::unbounded_channel();
903        handler.send_tx = Some(send_tx);
904        let (first_tx, first_rx) = tokio::sync::oneshot::channel();
905        let (second_tx, second_rx) = tokio::sync::oneshot::channel();
906        handler.enqueue_request("first", secret_json(json!({})), 0, first_tx);
907        handler.enqueue_request("second", secret_json(json!({})), 0, second_tx);
908
909        handler.fail_uncorrelated_error(crate::http::models::JsonRpcError {
910            code: -32600,
911            message: "Invalid Request".to_string(),
912            data: None,
913        });
914
915        assert_eq!(handler.pending.len(), 2);
916        handler.fail_pending("test cleanup");
917        assert!(first_rx.await.unwrap().is_err());
918        assert!(second_rx.await.unwrap().is_err());
919    }
920}
nautilus_derive/websocket/handler.rs

nautilus_derive/websocket/
handler.rs
