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_common::live::dst::task;
36use nautilus_core::string::secret::{SecretString, zeroize_json_value};
37use nautilus_live::task::{SharedTaskSlot, TaskJoinOutcome};
38use nautilus_network::{
39    RECONNECTED,
40    error::SendError,
41    websocket::{AuthTracker, WebSocketClient},
42};
43use serde_json::{Value, value::RawValue};
44use tokio_tungstenite::tungstenite::Message;
45use tokio_util::sync::CancellationToken;
46
47use super::{
48    client::UNAUTHENTICATED_CONNECTION_EPOCH,
49    error::DeriveWsError,
50    messages::{DeriveWsChannel, ParsedFrame, WsSubscribeParams, WsSubscriptionPayload},
51};
52use crate::{
53    book::sync::{BookSyncTracker, BookWrite},
54    http::models::JsonRpcRequest,
55};
56
57/// Outbound commands the outer client sends to the inner handler.
58#[derive(Debug)]
59pub(super) enum HandlerCommand {
60    /// Hand the active [`WebSocketClient`] to the handler.
61    SetClient(WebSocketClient),
62    /// Send a JSON-RPC request and resolve the oneshot when the venue replies.
63    /// `params` is pre-serialized JSON so the handler stays agnostic to the
64    /// per-method param types (login, subscribe, signed `private/*` bodies).
65    Request {
66        method: &'static str,
67        params: SecretString,
68        connection_epoch: Option<u64>,
69        cancellation: Option<CancellationToken>,
70        response_tx: tokio::sync::oneshot::Sender<Result<Box<RawValue>, DeriveWsError>>,
71    },
72    /// Gracefully tear down the WebSocket connection.
73    Disconnect,
74}
75
76/// Events emitted by the handler for the outer client and downstream consumers.
77#[derive(Debug, Clone)]
78pub enum DeriveWsMessage {
79    /// `public/login` succeeded. Consumed by the client's spawn loop to drive
80    /// resubscription; not forwarded to data/execution layers.
81    Authenticated,
82    /// Underlying transport reconnected; outer client triggers re-login and
83    /// resubscribes the tracked channels.
84    Reconnected,
85    /// Re-login or subscription replay exhausted its retry budget.
86    SessionRecoveryFailed(String),
87    /// Channel update pushed by the venue.
88    Subscription(WsSubscriptionPayload),
89    /// Book update admitted at ingress with its subscription generation.
90    BookSubscription {
91        /// Channel update pushed by the venue.
92        payload: WsSubscriptionPayload,
93        /// Cancelled when replacement, unsubscribe, or reconnect obsoletes the frame.
94        generation: CancellationToken,
95    },
96}
97
98#[derive(Debug)]
99struct SendCommand {
100    id: u64,
101    token: u64,
102    connection_epoch: u64,
103    payload: SecretString,
104    cancellation: Option<CancellationToken>,
105    book_generations: Vec<CancellationToken>,
106}
107
108#[derive(Debug)]
109struct SendFailure {
110    id: u64,
111    token: u64,
112    error: SendError,
113}
114
115#[derive(Debug)]
116struct PendingRequest {
117    token: u64,
118    cancellation: Option<CancellationToken>,
119    book_writes: Vec<(String, BookWrite)>,
120    response_tx: tokio::sync::oneshot::Sender<Result<Box<RawValue>, DeriveWsError>>,
121}
122
123/// Inner I/O loop. Lives in a Tokio task spawned by
124/// [`super::client::DeriveWebSocketClient::connect`].
125pub(super) struct FeedHandler {
126    signal: Arc<AtomicBool>,
127    book_sync: Option<(BookSyncTracker, Duration)>,
128    client: Option<Arc<WebSocketClient>>,
129    cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
130    cmd_closed: bool,
131    raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
132    raw_closed: bool,
133    next_id: Arc<AtomicU64>,
134    next_send_token: u64,
135    pending: AHashMap<u64, PendingRequest>,
136    send_tx: Option<tokio::sync::mpsc::UnboundedSender<SendCommand>>,
137    send_failure_rx: tokio::sync::mpsc::UnboundedReceiver<SendFailure>,
138    send_task: Arc<SharedTaskSlot<()>>,
139    auth_tracker: AuthTracker,
140    authenticated_epoch: Arc<AtomicU64>,
141}
142
143impl FeedHandler {
144    pub(super) fn new_with_send_task(
145        signal: Arc<AtomicBool>,
146        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
147        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
148        next_id: Arc<AtomicU64>,
149        auth_tracker: AuthTracker,
150        authenticated_epoch: Arc<AtomicU64>,
151        send_task: Arc<SharedTaskSlot<()>>,
152    ) -> Self {
153        let (_, send_failure_rx) = tokio::sync::mpsc::unbounded_channel();
154
155        Self {
156            signal,
157            book_sync: None,
158            client: None,
159            cmd_rx,
160            cmd_closed: false,
161            raw_rx,
162            raw_closed: false,
163            next_id,
164            next_send_token: 0,
165            pending: AHashMap::new(),
166            send_tx: None,
167            send_failure_rx,
168            send_task,
169            auth_tracker,
170            authenticated_epoch,
171        }
172    }
173
174    pub(super) fn with_book_sync(mut self, book_sync: Option<(BookSyncTracker, Duration)>) -> Self {
175        self.book_sync = book_sync;
176        self
177    }
178
179    /// Drains the next event from the underlying channels, processes it, and
180    /// returns the resulting outbound message (if any). Returns `None` when
181    /// the handler is shutting down or both channels closed.
182    pub(super) async fn next(&mut self) -> Option<DeriveWsMessage> {
183        let mut command_ready = true;
184
185        loop {
186            tokio::select! {
187                biased;
188                Some(failure) = self.send_failure_rx.recv() => {
189                    self.handle_send_failure(&failure);
190                }
191                cmd = self.cmd_rx.recv(), if !self.cmd_closed && command_ready => {
192                    command_ready = false;
193
194                    match cmd {
195                        None => {
196                            self.cmd_closed = true;
197
198                            if self.raw_closed {
199                                self.shutdown_send_path(
200                                    "WebSocket handler stopped before response was received",
201                                )
202                                .await;
203                                return None;
204                            }
205                        }
206                        Some(HandlerCommand::SetClient(client)) => {
207                            log::debug!("Setting WebSocket client in Derive handler");
208                            let client = Arc::new(client);
209                            self.stop_send_worker().await;
210                            self.start_send_worker(Arc::clone(&client));
211                            self.client = Some(client);
212                        }
213                        Some(HandlerCommand::Request {
214                            method,
215                            params,
216                            connection_epoch,
217                            cancellation,
218                            response_tx,
219                        }) => {
220                            self.dispatch_request(method, params, connection_epoch, cancellation, response_tx);
221                        }
222                        Some(HandlerCommand::Disconnect) => {
223                            log::debug!("Derive handler received disconnect command");
224                            self.shutdown_send_path(
225                                "WebSocket disconnected before response was received",
226                            )
227                            .await;
228
229                            if let Some(ref client) = self.client {
230                                client.disconnect().await;
231                            }
232                            self.signal.store(true, Ordering::SeqCst);
233                            return None;
234                        }
235                    }
236                }
237
238                raw = self.raw_rx.recv(), if !self.raw_closed => {
239                    command_ready = true;
240
241                    match raw {
242                        None => {
243                            self.raw_closed = true;
244
245                            if self.cmd_closed {
246                                self.shutdown_send_path(
247                                    "WebSocket handler stopped before response was received",
248                                )
249                                .await;
250                                return None;
251                            }
252                        }
253                        Some(Message::Text(text)) => {
254                            if text.as_str() == RECONNECTED {
255                                log::info!("Derive WebSocket reconnected sentinel received");
256                                self.auth_tracker.invalidate();
257                                self.authenticated_epoch.store(
258                                    UNAUTHENTICATED_CONNECTION_EPOCH,
259                                    Ordering::Release,
260                                );
261                                self.restart_send_worker(
262                                    "WebSocket reconnected before response was received",
263                                )
264                                .await;
265
266                                if let Some((tracker, timeout)) = &self.book_sync {
267                                    tracker.reset_on_reconnect(*timeout);
268                                }
269                                return Some(DeriveWsMessage::Reconnected);
270                            }
271
272                            match ParsedFrame::parse(&text) {
273                                Ok(ParsedFrame::Response { id, result, error }) => {
274                                    if let Some(pending) = self.pending.remove(&id) {
275                                        let outcome = match (result, error) {
276                                            (_, Some(err)) => Err(DeriveWsError::JsonRpc {
277                                                code: err.code,
278                                                message: err.message,
279                                                data: err.data,
280                                            }),
281                                            (Some(value), None) => Ok(value),
282                                            (None, None) => Ok(RawValue::NULL.to_owned()),
283                                        };
284
285                                        if let Ok(value) = &outcome
286                                            && let Some((tracker, _)) = &self.book_sync {
287                                            open_book_gates(&pending, value, tracker);
288                                        }
289                                        let _ = pending.response_tx.send(outcome);
290                                    } else {
291                                        log::debug!(
292                                            "Derive WebSocket response with unknown id={id} dropped",
293                                        );
294                                    }
295                                }
296                                Ok(ParsedFrame::Subscription(payload)) => {
297                                    if let Some((tracker, _)) = &self.book_sync
298                                        && payload.channel.as_str().starts_with("orderbook.") {
299                                        let Some(write) = tracker.write(payload.channel.as_str()) else {
300                                            continue;
301                                        };
302
303                                        if write.generation.is_cancelled() || write.gate.lock().is_closed() {
304                                            continue;
305                                        }
306                                        return Some(DeriveWsMessage::BookSubscription {
307                                            payload,
308                                            generation: write.generation,
309                                        });
310                                    }
311                                    return Some(DeriveWsMessage::Subscription(payload));
312                                }
313                                Ok(ParsedFrame::UncorrelatedError(error)) => {
314                                    self.fail_uncorrelated_error(error);
315                                }
316                                Ok(ParsedFrame::Unknown(_)) => {
317                                    log::debug!("Derive WebSocket unknown frame, bytes={}", text.len());
318                                }
319                                Err(e) => {
320                                    log::error!(
321                                        "Derive WebSocket frame parse error: {:?}, line={}, column={}, bytes={}",
322                                        e.classify(),
323                                        e.line(),
324                                        e.column(),
325                                        text.len(),
326                                    );
327                                }
328                            }
329                        }
330                        Some(Message::Ping(data)) => {
331                            if let Some(ref client) = self.client
332                                && let Err(e) = client.send_pong(data.to_vec()).await {
333                                log::error!("Derive WebSocket send_pong failed: {e}");
334                            }
335                        }
336                        Some(Message::Close(_)) => {
337                            log::debug!("Derive WebSocket close frame received");
338                            self.shutdown_send_path(
339                                "WebSocket closed before response was received",
340                            )
341                            .await;
342                            return None;
343                        }
344                        Some(_) => {}
345                    }
346                }
347
348                () = task::yield_now(), if !command_ready => command_ready = true,
349
350                else => {
351                    log::debug!("Derive handler shutting down: channels closed");
352                    self.shutdown_send_path(
353                        "WebSocket handler stopped before response was received",
354                    )
355                    .await;
356                    return None;
357                }
358            }
359        }
360    }
361
362    fn dispatch_request(
363        &mut self,
364        method: &'static str,
365        params: SecretString,
366        connection_epoch: Option<u64>,
367        cancellation: Option<CancellationToken>,
368        response_tx: tokio::sync::oneshot::Sender<Result<Box<RawValue>, DeriveWsError>>,
369    ) {
370        if cancellation
371            .as_ref()
372            .is_some_and(CancellationToken::is_cancelled)
373        {
374            let _ = response_tx.send(Err(DeriveWsError::NotConnected));
375            return;
376        }
377
378        let Some(client) = self.client.as_ref() else {
379            let _ = response_tx.send(Err(DeriveWsError::NotConnected));
380            return;
381        };
382
383        let current_epoch = client.connection_epoch();
384
385        let connection_epoch = if method.starts_with("private/") {
386            let authenticated_epoch = self.authenticated_epoch.load(Ordering::Acquire);
387            if !client.connection_mode().is_active() || authenticated_epoch != current_epoch {
388                let _ = response_tx.send(Err(DeriveWsError::Authentication {
389                    operation: method.to_string(),
390                    reason: "WebSocket session is not authenticated".to_string(),
391                }));
392
393                return;
394            }
395
396            authenticated_epoch
397        } else {
398            connection_epoch.unwrap_or(current_epoch)
399        };
400
401        if !client.connection_mode().is_active() || connection_epoch != current_epoch {
402            let _ = response_tx.send(Err(DeriveWsError::NotConnected));
403            return;
404        }
405
406        self.enqueue_request(method, params, connection_epoch, cancellation, response_tx);
407    }
408
409    fn enqueue_request(
410        &mut self,
411        method: &'static str,
412        params_json: SecretString,
413        connection_epoch: u64,
414        cancellation: Option<CancellationToken>,
415        response_tx: tokio::sync::oneshot::Sender<Result<Box<RawValue>, DeriveWsError>>,
416    ) {
417        let Some(send_tx) = self.send_tx.clone() else {
418            let _ = response_tx.send(Err(DeriveWsError::NotConnected));
419            return;
420        };
421
422        let id = self.next_id.fetch_add(1, Ordering::Relaxed);
423        let token = self.next_send_token;
424        self.next_send_token = self.next_send_token.wrapping_add(1);
425
426        let params: Value = match serde_json::from_str(params_json.expose_secret()) {
427            Ok(params) => params,
428            Err(e) => {
429                let _ = response_tx.send(Err(DeriveWsError::Serde(e)));
430                return;
431            }
432        };
433
434        drop(params_json);
435
436        let book_writes = self
437            .book_sync
438            .as_ref()
439            .map_or_else(Vec::new, |(tracker, _)| {
440                params
441                    .get("channels")
442                    .and_then(Value::as_array)
443                    .into_iter()
444                    .flatten()
445                    .filter_map(Value::as_str)
446                    .filter_map(|channel| {
447                        tracker
448                            .write(channel)
449                            .map(|write| (channel.to_owned(), write))
450                    })
451                    .collect::<Vec<_>>()
452            });
453
454        let book_generations = if cancellation.is_some() {
455            book_writes
456                .iter()
457                .map(|(_, write)| write.generation.clone())
458                .collect()
459        } else {
460            Vec::new()
461        };
462
463        let book_writes = if method == "subscribe" {
464            book_writes
465        } else {
466            Vec::new()
467        };
468
469        let mut request = JsonRpcRequest::new(id, method, params);
470        let serialized = serde_json::to_string(&request);
471        zeroize_json_value(&mut request.params);
472
473        let payload = match serialized {
474            Ok(payload) => SecretString::from(payload),
475            Err(e) => {
476                let _ = response_tx.send(Err(DeriveWsError::Serde(e)));
477                return;
478            }
479        };
480
481        self.pending.insert(
482            id,
483            PendingRequest {
484                token,
485                cancellation: cancellation.clone(),
486                book_writes,
487                response_tx,
488            },
489        );
490
491        log::debug!("Derive WebSocket sending `{method}` id={id}");
492
493        if send_tx
494            .send(SendCommand {
495                id,
496                token,
497                connection_epoch,
498                payload,
499                cancellation,
500                book_generations,
501            })
502            .is_err()
503            && self
504                .pending
505                .get(&id)
506                .is_some_and(|pending| pending.token == token)
507            && let Some(pending) = self.pending.remove(&id)
508        {
509            let _ = pending.response_tx.send(Err(DeriveWsError::NotConnected));
510        }
511    }
512
513    fn start_send_worker(&mut self, client: Arc<WebSocketClient>) {
514        if !self.send_task.is_empty() {
515            log::error!("Cannot start Derive WebSocket send worker while the prior task is owned");
516            return;
517        }
518
519        let (send_tx, send_rx) = tokio::sync::mpsc::unbounded_channel();
520        let (failure_tx, failure_rx) = tokio::sync::mpsc::unbounded_channel();
521        self.send_tx = Some(send_tx);
522        self.send_failure_rx = failure_rx;
523
524        if let Err(e) = self
525            .send_task
526            .spawn(run_send_worker(client, send_rx, failure_tx))
527        {
528            self.send_tx.take();
529            log::error!("Failed to start Derive WebSocket send worker: {e}");
530        }
531    }
532
533    async fn stop_send_worker(&mut self) {
534        self.send_tx.take();
535
536        if let Some(outcome) = self
537            .send_task
538            .finish(Duration::from_secs(1), Duration::from_secs(2))
539            .await
540        {
541            match outcome {
542                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
543                TaskJoinOutcome::Failed(e) => {
544                    log::error!("Derive WebSocket send worker failed: {e}");
545                }
546                TaskJoinOutcome::Incomplete => {
547                    log::error!("Derive WebSocket send worker did not stop after abort");
548                }
549            }
550        }
551    }
552
553    async fn restart_send_worker(&mut self, reason: &str) {
554        self.stop_send_worker().await;
555        self.fail_pending(reason);
556
557        if let Some(client) = self.client.clone() {
558            self.start_send_worker(client);
559        }
560    }
561
562    async fn shutdown_send_path(&mut self, reason: &str) {
563        self.stop_send_worker().await;
564        self.fail_pending(reason);
565    }
566
567    fn handle_send_failure(&mut self, failure: &SendFailure) {
568        if !self
569            .pending
570            .get(&failure.id)
571            .is_some_and(|pending| pending.token == failure.token)
572        {
573            return;
574        }
575
576        if let Some(pending) = self.pending.remove(&failure.id) {
577            let _ = pending
578                .response_tx
579                .send(Err(send_error_to_derive(&failure.error)));
580        }
581    }
582
583    fn fail_uncorrelated_error(&mut self, error: crate::http::models::JsonRpcError) {
584        if self.pending.len() != 1 {
585            log::warn!(
586                "Derive WebSocket uncorrelated JSON-RPC error with {} pending requests: code={}, message={}",
587                self.pending.len(),
588                error.code,
589                error.message,
590            );
591            return;
592        }
593
594        if let Some((_, pending)) = self.pending.drain().next() {
595            let _ = pending.response_tx.send(Err(DeriveWsError::JsonRpc {
596                code: error.code,
597                message: error.message,
598                data: error.data,
599            }));
600        }
601    }
602
603    fn fail_pending(&mut self, reason: &str) {
604        if self.pending.is_empty() {
605            return;
606        }
607
608        log::debug!(
609            "Failing {} pending Derive WebSocket request(s): {reason}",
610            self.pending.len(),
611        );
612
613        for (_, pending) in self.pending.drain() {
614            let _ = pending
615                .response_tx
616                .send(Err(DeriveWsError::transport(reason.to_string())));
617        }
618    }
619}
620
621impl Drop for FeedHandler {
622    fn drop(&mut self) {
623        self.send_tx.take();
624        self.send_task.abort();
625    }
626}
627
628async fn run_send_worker(
629    client: Arc<WebSocketClient>,
630    mut send_rx: tokio::sync::mpsc::UnboundedReceiver<SendCommand>,
631    failure_tx: tokio::sync::mpsc::UnboundedSender<SendFailure>,
632) {
633    while let Some(command) = send_rx.recv().await {
634        let result = tokio::select! {
635            biased;
636            () = async {
637                if let Some(cancel) = &command.cancellation {
638                    cancel.cancelled().await;
639                } else {
640                    std::future::pending::<()>().await;
641                }
642            } => Err(SendError::Closed),
643            result = async {
644                if command.book_generations.iter().any(CancellationToken::is_cancelled) {
645                    return Err(SendError::Closed);
646                }
647                client.send_text_on_connection(
648                    command.payload.expose_secret().to_owned(),
649                    None,
650                    command.connection_epoch,
651                ).await
652            } => result,
653        };
654
655        if let Err(e) = result
656            && failure_tx
657                .send(SendFailure {
658                    id: command.id,
659                    token: command.token,
660                    error: e,
661                })
662                .is_err()
663        {
664            return;
665        }
666    }
667}
668
669fn open_book_gates(pending: &PendingRequest, value: &RawValue, tracker: &BookSyncTracker) {
670    let is_pending = || {
671        !pending.response_tx.is_closed()
672            && !pending
673                .cancellation
674                .as_ref()
675                .is_some_and(CancellationToken::is_cancelled)
676    };
677
678    if pending.book_writes.is_empty() || !is_pending() {
679        return;
680    }
681
682    let Ok(result) = serde_json::from_str::<super::messages::WsSubscribeResult>(value.get()) else {
683        return;
684    };
685
686    for (channel, write) in &pending.book_writes {
687        let topic = DeriveWsChannel::from(channel.clone());
688
689        if result
690            .status
691            .get(&topic)
692            .is_some_and(|status| matches!(status.as_str(), "ok" | "already subscribed"))
693            && result.channels.contains(&topic)
694        {
695            tracker.confirm_subscription(write, is_pending);
696        }
697    }
698}
699
700/// Builds `subscribe` params from a single channel topic.
701#[must_use]
702pub(super) fn subscribe_params(channel: DeriveWsChannel) -> WsSubscribeParams {
703    WsSubscribeParams {
704        channels: vec![channel],
705    }
706}
707
708fn send_error_to_derive(error: &SendError) -> DeriveWsError {
709    match error {
710        SendError::InvalidInput(_)
711        | SendError::BufferFull
712        | SendError::Closed
713        | SendError::Timeout
714        | SendError::ConnectionChanged => {
715            log::warn!("Derive WebSocket request was not sent: {error}");
716            DeriveWsError::NotConnected
717        }
718        SendError::BrokenPipe(_) | SendError::WriteTimeout => {
719            DeriveWsError::transport(error.to_string())
720        }
721    }
722}
723
724/// Convenience wrapper that produces the `subscribe` params for the
725/// `ticker_slim.{instrument_name}.{interval}` channel.
726#[must_use]
727pub(super) fn ticker_subscribe_params(instrument_name: &str, interval: &str) -> WsSubscribeParams {
728    subscribe_params(DeriveWsChannel::ticker_slim(instrument_name, interval))
729}
730
731/// Convenience wrapper that produces the `subscribe` params for the
732/// `orderbook.{instrument_name}.{group}.{depth}` channel.
733#[must_use]
734pub(super) fn orderbook_subscribe_params(
735    instrument_name: &str,
736    group: &str,
737    depth: &str,
738) -> WsSubscribeParams {
739    subscribe_params(DeriveWsChannel::orderbook(instrument_name, group, depth))
740}
741
742/// Convenience wrapper that produces the `subscribe` params for the
743/// `trades.{instrument_type}.{currency}` channel.
744#[must_use]
745pub(super) fn trades_subscribe_params(instrument_type: &str, currency: &str) -> WsSubscribeParams {
746    subscribe_params(DeriveWsChannel::trades(instrument_type, currency))
747}
748
749#[cfg(test)]
750mod tests {
751    use futures_util::StreamExt;
752    use nautilus_model::identifiers::InstrumentId;
753    use nautilus_network::websocket::WebSocketConfig;
754    use rstest::rstest;
755    use serde_json::json;
756
757    use super::*;
758
759    fn unauthenticated_epoch() -> Arc<AtomicU64> {
760        Arc::new(AtomicU64::new(u64::MAX))
761    }
762
763    fn secret_json(mut value: Value) -> SecretString {
764        let serialized = SecretString::from(value.to_string());
765        zeroize_json_value(&mut value);
766        serialized
767    }
768
769    #[rstest]
770    #[case(SendError::InvalidInput("invalid".to_string()), false)]
771    #[case(SendError::BufferFull, false)]
772    #[case(SendError::Closed, false)]
773    #[case(SendError::Timeout, false)]
774    #[case(SendError::ConnectionChanged, false)]
775    #[case(SendError::BrokenPipe("lost connection".to_string()), true)]
776    #[case(SendError::WriteTimeout, true)]
777    fn test_send_error_preserves_write_ambiguity(
778        #[case] error: SendError,
779        #[case] ambiguous: bool,
780    ) {
781        let reason = error.to_string();
782        let error = send_error_to_derive(&error);
783        assert_eq!(
784            crate::common::retry::is_write_outcome_ambiguous_ws(&error),
785            ambiguous
786        );
787
788        if ambiguous {
789            let DeriveWsError::Transport(message) = error else {
790                panic!("expected ambiguous transport error: {error:?}");
791            };
792
793            assert_eq!(message, reason);
794        } else {
795            assert!(matches!(error, DeriveWsError::NotConnected));
796        }
797    }
798
799    #[rstest]
800    fn test_queued_request_debug_redacts_payload() {
801        let secret = "signature-secret";
802        let (response_tx, _response_rx) = tokio::sync::oneshot::channel();
803
804        let command = HandlerCommand::Request {
805            method: "public/login",
806            params: SecretString::from(format!(r#"{{"signature":"{secret}"}}"#)),
807            connection_epoch: None,
808            cancellation: None,
809            response_tx,
810        };
811
812        let send = SendCommand {
813            cancellation: None,
814            book_generations: Vec::new(),
815            id: 1,
816            token: 1,
817            connection_epoch: 1,
818            payload: SecretString::from(format!(r#"{{"signature":"{secret}"}}"#)),
819        };
820
821        let debug = format!("{command:?} {send:?}");
822
823        assert!(!debug.contains(secret));
824    }
825
826    #[rstest]
827    #[case::episode(true)]
828    #[case::generation(false)]
829    #[tokio::test]
830    async fn cancelled_book_writes_never_reach_socket(#[case] episode: bool) {
831        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
832        let address = listener.local_addr().unwrap();
833        let peer = tokio::spawn(async move {
834            let (socket, _) = listener.accept().await.unwrap();
835            let mut socket = tokio_tungstenite::accept_async(socket).await.unwrap();
836            let mut messages = Vec::new();
837
838            while let Some(Ok(message)) = socket.next().await {
839                match message {
840                    Message::Text(text) => messages.push(text.to_string()),
841                    Message::Close(_) => break,
842                    _ => {}
843                }
844            }
845
846            messages
847        });
848
849        let config = WebSocketConfig::builder()
850            .url(format!("ws://{address}"))
851            .build()
852            .unwrap();
853
854        let client = Arc::new(
855            WebSocketClient::builder()
856                .config(config)
857                .message_handler(Arc::new(|_| {}))
858                .connect()
859                .await
860                .unwrap(),
861        );
862        let cancel = CancellationToken::new();
863        let generation = CancellationToken::new();
864
865        if episode {
866            cancel.cancel();
867        } else {
868            generation.cancel();
869        }
870
871        let (send_tx, send_rx) = tokio::sync::mpsc::unbounded_channel();
872        let (failure_tx, mut failure_rx) = tokio::sync::mpsc::unbounded_channel();
873        send_tx
874            .send(SendCommand {
875                id: 17,
876                token: 29,
877                connection_epoch: client.connection_epoch(),
878                payload: SecretString::from("cancelled"),
879                cancellation: Some(cancel),
880                book_generations: vec![generation],
881            })
882            .unwrap();
883
884        send_tx
885            .send(SendCommand {
886                id: 31,
887                token: 43,
888                connection_epoch: client.connection_epoch(),
889                payload: SecretString::from("current"),
890                cancellation: None,
891                book_generations: Vec::new(),
892            })
893            .unwrap();
894
895        drop(send_tx);
896
897        run_send_worker(Arc::clone(&client), send_rx, failure_tx).await;
898        let failure = failure_rx.try_recv().unwrap();
899        client.disconnect().await;
900        let messages = tokio::time::timeout(std::time::Duration::from_secs(3), peer)
901            .await
902            .unwrap()
903            .unwrap();
904
905        assert_eq!(failure.id, 17);
906        assert_eq!(failure.token, 29);
907        assert!(matches!(failure.error, SendError::Closed));
908        assert!(failure_rx.try_recv().is_err());
909        assert_eq!(messages, ["current"]);
910    }
911
912    #[rstest]
913    #[case::tracked(true)]
914    #[case::legacy(false)]
915    #[tokio::test]
916    async fn unowned_book_frames_are_dropped_only_on_tracked_streams(#[case] tracked: bool) {
917        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
918        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
919        let mut handler = feed_handler(
920            Arc::new(AtomicBool::new(false)),
921            cmd_rx,
922            raw_rx,
923            Arc::new(AtomicU64::new(1)),
924            AuthTracker::default(),
925            unauthenticated_epoch(),
926        )
927        .with_book_sync(tracked.then(|| (BookSyncTracker::default(), Duration::from_secs(10))));
928        let data: Value =
929            serde_json::from_str(include_str!("../../test_data/spot/ws_orderbook_eth.json"))
930                .unwrap();
931        let channel = "orderbook.ETH-USDC.1.10";
932        raw_tx
933            .send(Message::Text(
934                json!({"method": "subscription", "params": {"channel": channel, "data": data}})
935                    .to_string()
936                    .into(),
937            ))
938            .unwrap();
939        raw_tx.send(Message::Text(RECONNECTED.into())).unwrap();
940        let event = tokio::time::timeout(Duration::from_secs(1), handler.next())
941            .await
942            .unwrap()
943            .unwrap();
944
945        if tracked {
946            assert!(matches!(event, DeriveWsMessage::Reconnected));
947        } else {
948            let DeriveWsMessage::Subscription(payload) = event else {
949                panic!("expected legacy book event");
950            };
951
952            assert_eq!(payload.channel.as_str(), channel);
953            assert_eq!(
954                serde_json::from_str::<Value>(payload.data.get()).unwrap(),
955                data
956            );
957        }
958    }
959
960    #[rstest]
961    #[tokio::test]
962    async fn book_ack_opens_gate_at_ingress_and_drops_pre_ack_frames() {
963        let tracker = BookSyncTracker::default();
964        let channel = "orderbook.ETH-USDC.1.10";
965        tracker.record_subscription(channel, InstrumentId::from("ETH-USDC.DERIVE"));
966        let write = tracker.write(channel).unwrap();
967        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
968        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
969        let mut handler = feed_handler(
970            Arc::new(AtomicBool::new(false)),
971            cmd_rx,
972            raw_rx,
973            Arc::new(AtomicU64::new(1)),
974            AuthTracker::default(),
975            Arc::new(AtomicU64::new(UNAUTHENTICATED_CONNECTION_EPOCH)),
976        )
977        .with_book_sync(Some((tracker, Duration::from_secs(10))));
978        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
979        handler.pending.insert(
980            1,
981            PendingRequest {
982                token: 1,
983                cancellation: Some(write.generation.clone()),
984                book_writes: vec![(channel.to_owned(), write.clone())],
985                response_tx,
986            },
987        );
988
989        let ack: Value =
990            serde_json::from_str(include_str!("../../test_data/spot/ws_subscribe_ack.json"))
991                .unwrap();
992        let mut data: Value =
993            serde_json::from_str(include_str!("../../test_data/spot/ws_orderbook_eth.json"))
994                .unwrap();
995
996        let notification = |data: &Value| {
997            Message::Text(
998                json!({"method": "subscription", "params": {"channel": channel, "data": data}})
999                    .to_string()
1000                    .into(),
1001            )
1002        };
1003
1004        raw_tx.send(notification(&data)).unwrap();
1005        raw_tx
1006            .send(Message::Text(
1007                json!({"id": 1, "result": ack["result"]}).to_string().into(),
1008            ))
1009            .unwrap();
1010        data["timestamp"] = json!(data["timestamp"].as_i64().unwrap() + 1);
1011        raw_tx.send(notification(&data)).unwrap();
1012        let event = tokio::time::timeout(Duration::from_secs(1), handler.next())
1013            .await
1014            .unwrap()
1015            .unwrap();
1016
1017        let DeriveWsMessage::BookSubscription {
1018            payload,
1019            generation,
1020        } = event
1021        else {
1022            panic!("expected gated book event");
1023        };
1024
1025        assert_eq!(
1026            serde_json::from_str::<Value>(payload.data.get()).unwrap(),
1027            data
1028        );
1029        assert!(!generation.is_cancelled());
1030        assert!(!write.gate.lock().is_closed());
1031        assert_eq!(
1032            serde_json::from_str::<Value>(response_rx.await.unwrap().unwrap().get()).unwrap(),
1033            ack["result"]
1034        );
1035    }
1036
1037    #[rstest]
1038    #[case::cancelled(true, false, false, true)]
1039    #[case::obsolete(false, true, false, true)]
1040    #[case::abandoned(false, false, true, true)]
1041    #[case::current(false, false, false, false)]
1042    fn book_ack_opens_only_live_request(
1043        #[case] cancelled: bool,
1044        #[case] obsolete: bool,
1045        #[case] abandoned: bool,
1046        #[case] expected_closed: bool,
1047    ) {
1048        let tracker = BookSyncTracker::default();
1049        let channel = "orderbook.ETH-USDC.1.10";
1050        tracker.record_subscription(channel, InstrumentId::from("ETH-USDC.DERIVE"));
1051        let write = tracker.write(channel).unwrap();
1052        let gate = write.gate.clone();
1053        let generation = write.generation.clone();
1054        let cancel = CancellationToken::new();
1055        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1056        let _response_rx = if abandoned {
1057            drop(response_rx);
1058            None
1059        } else {
1060            Some(response_rx)
1061        };
1062
1063        let pending = PendingRequest {
1064            token: 1,
1065            cancellation: Some(cancel.clone()),
1066            book_writes: vec![(channel.to_owned(), write)],
1067            response_tx,
1068        };
1069
1070        if cancelled {
1071            cancel.cancel();
1072        }
1073
1074        if obsolete {
1075            generation.cancel();
1076        }
1077
1078        let ack: Value =
1079            serde_json::from_str(include_str!("../../test_data/spot/ws_subscribe_ack.json"))
1080                .unwrap();
1081        let result = RawValue::from_string(ack["result"].to_string()).unwrap();
1082        open_book_gates(&pending, &result, &tracker);
1083        assert_eq!(gate.lock().is_closed(), expected_closed);
1084    }
1085
1086    mod dst {
1087        use super::*;
1088
1089        #[cfg_attr(all(feature = "simulation", madsim), madsim::test)]
1090        #[cfg_attr(not(all(feature = "simulation", madsim)), tokio::test)]
1091        async fn commands_and_frames_progress_when_both_queues_are_ready() {
1092            let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1093            let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1094            let mut handler = feed_handler(
1095                Arc::new(AtomicBool::new(false)),
1096                cmd_rx,
1097                raw_rx,
1098                Arc::new(AtomicU64::new(1)),
1099                AuthTracker::new(),
1100                unauthenticated_epoch(),
1101            );
1102            let mut responses = Vec::new();
1103
1104            for _ in 0..2 {
1105                let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1106                cmd_tx
1107                    .send(HandlerCommand::Request {
1108                        method: "subscribe",
1109                        params: secret_json(json!({"channels": ["trades.perp.ETH"]})),
1110                        connection_epoch: None,
1111                        cancellation: None,
1112                        response_tx,
1113                    })
1114                    .unwrap();
1115
1116                responses.push(response_rx);
1117            }
1118
1119            raw_tx
1120                .send(Message::text(
1121                    r#"{"method":"subscription","params":{"channel":"trades.perp.ETH","data":[]}}"#,
1122                ))
1123                .unwrap();
1124            let event =
1125                nautilus_common::live::dst::time::timeout(Duration::from_secs(1), handler.next())
1126                    .await
1127                    .unwrap()
1128                    .unwrap();
1129
1130            let DeriveWsMessage::Subscription(payload) = event else {
1131                panic!("expected subscription frame");
1132            };
1133
1134            assert_eq!(payload.channel.as_str(), "trades.perp.ETH");
1135            assert_eq!(payload.data.get(), "[]");
1136            assert!(matches!(
1137                responses[0].try_recv().unwrap(),
1138                Err(DeriveWsError::NotConnected)
1139            ));
1140            assert_eq!(
1141                responses[1].try_recv().unwrap_err(),
1142                tokio::sync::oneshot::error::TryRecvError::Empty
1143            );
1144            cmd_tx.send(HandlerCommand::Disconnect).unwrap();
1145            assert!(
1146                nautilus_common::live::dst::time::timeout(Duration::from_secs(1), handler.next())
1147                    .await
1148                    .unwrap()
1149                    .is_none()
1150            );
1151            assert!(matches!(
1152                responses[1].try_recv().unwrap(),
1153                Err(DeriveWsError::NotConnected)
1154            ));
1155        }
1156    }
1157
1158    fn feed_handler(
1159        signal: Arc<AtomicBool>,
1160        cmd_rx: tokio::sync::mpsc::UnboundedReceiver<HandlerCommand>,
1161        raw_rx: tokio::sync::mpsc::UnboundedReceiver<Message>,
1162        next_id: Arc<AtomicU64>,
1163        auth_tracker: AuthTracker,
1164        authenticated_epoch: Arc<AtomicU64>,
1165    ) -> FeedHandler {
1166        FeedHandler::new_with_send_task(
1167            signal,
1168            cmd_rx,
1169            raw_rx,
1170            next_id,
1171            auth_tracker,
1172            authenticated_epoch,
1173            Arc::new(SharedTaskSlot::new()),
1174        )
1175    }
1176
1177    #[rstest]
1178    fn test_subscribe_params_carries_single_channel() {
1179        let params = subscribe_params(DeriveWsChannel::ticker_slim("ETH-PERP", "1000"));
1180        assert_eq!(
1181            params.channels,
1182            vec![DeriveWsChannel::ticker_slim("ETH-PERP", "1000")],
1183        );
1184    }
1185
1186    #[rstest]
1187    fn test_ticker_subscribe_params_formats_topic() {
1188        let params = ticker_subscribe_params("ETH-PERP", "1000");
1189        assert_eq!(
1190            params.channels,
1191            vec![DeriveWsChannel::ticker_slim("ETH-PERP", "1000")],
1192        );
1193    }
1194
1195    #[rstest]
1196    fn test_orderbook_subscribe_params_formats_topic() {
1197        let params = orderbook_subscribe_params("ETH-PERP", "1", "10");
1198        assert_eq!(
1199            params.channels,
1200            vec![DeriveWsChannel::orderbook("ETH-PERP", "1", "10")],
1201        );
1202    }
1203
1204    #[rstest]
1205    fn test_trades_subscribe_params_formats_topic() {
1206        let params = trades_subscribe_params("perp", "ETH");
1207        assert_eq!(
1208            params.channels,
1209            vec![DeriveWsChannel::trades("perp", "ETH")],
1210        );
1211    }
1212
1213    #[rstest]
1214    #[tokio::test]
1215    async fn test_dispatch_request_without_client_returns_not_connected() {
1216        // Requests issued before SetClient must fail fast rather than hang.
1217        let signal = Arc::new(AtomicBool::new(false));
1218        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1219        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1220        let next_id = Arc::new(AtomicU64::new(1));
1221        let auth_tracker = AuthTracker::new();
1222        let mut handler = feed_handler(
1223            signal,
1224            cmd_rx,
1225            raw_rx,
1226            next_id,
1227            auth_tracker,
1228            unauthenticated_epoch(),
1229        );
1230
1231        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1232        let params =
1233            secret_json(serde_json::to_value(WsSubscribeParams { channels: vec![] }).unwrap());
1234        handler.dispatch_request("public/login", params, None, None, response_tx);
1235
1236        let outcome = response_rx.await.expect("oneshot resolved");
1237        match outcome {
1238            Err(DeriveWsError::NotConnected) => {}
1239            other => panic!("expected NotConnected, was {other:?}"),
1240        }
1241    }
1242
1243    #[rstest]
1244    #[tokio::test]
1245    async fn test_dispatch_registers_pending_requests_before_ordered_queueing() {
1246        let signal = Arc::new(AtomicBool::new(false));
1247        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1248        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1249        let next_id = Arc::new(AtomicU64::new(1));
1250        let auth_tracker = AuthTracker::new();
1251        let mut handler = feed_handler(
1252            signal,
1253            cmd_rx,
1254            raw_rx,
1255            next_id,
1256            auth_tracker,
1257            unauthenticated_epoch(),
1258        );
1259        let (send_tx, mut send_rx) = tokio::sync::mpsc::unbounded_channel();
1260        handler.send_tx = Some(send_tx);
1261
1262        let (first_tx, _first_rx) = tokio::sync::oneshot::channel();
1263        let (second_tx, _second_rx) = tokio::sync::oneshot::channel();
1264        handler.enqueue_request(
1265            "first",
1266            secret_json(json!({"sequence": 1})),
1267            0,
1268            None,
1269            first_tx,
1270        );
1271        handler.enqueue_request(
1272            "second",
1273            secret_json(json!({"sequence": 2})),
1274            0,
1275            None,
1276            second_tx,
1277        );
1278
1279        let first = send_rx.recv().await.expect("first queued send");
1280        let second = send_rx.recv().await.expect("second queued send");
1281        let first_payload: Value = serde_json::from_str(first.payload.expose_secret()).unwrap();
1282        let second_payload: Value = serde_json::from_str(second.payload.expose_secret()).unwrap();
1283
1284        assert_eq!(first.id, 1);
1285        assert_eq!(second.id, 2);
1286        assert_eq!(first_payload["method"], "first");
1287        assert_eq!(second_payload["method"], "second");
1288        assert_eq!(handler.pending.get(&first.id).unwrap().token, first.token);
1289        assert_eq!(handler.pending.get(&second.id).unwrap().token, second.token);
1290    }
1291
1292    #[rstest]
1293    #[tokio::test]
1294    async fn test_late_send_failure_does_not_remove_reused_request_id() {
1295        let signal = Arc::new(AtomicBool::new(false));
1296        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1297        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1298        let next_id = Arc::new(AtomicU64::new(1));
1299        let auth_tracker = AuthTracker::new();
1300        let mut handler = feed_handler(
1301            signal,
1302            cmd_rx,
1303            raw_rx,
1304            Arc::clone(&next_id),
1305            auth_tracker,
1306            unauthenticated_epoch(),
1307        );
1308        let (send_tx, mut send_rx) = tokio::sync::mpsc::unbounded_channel();
1309        handler.send_tx = Some(send_tx);
1310
1311        let (old_tx, old_rx) = tokio::sync::oneshot::channel();
1312        handler.enqueue_request("first", secret_json(json!({})), 0, None, old_tx);
1313        let old_send = send_rx.recv().await.expect("old queued send");
1314        let old_pending = handler.pending.remove(&old_send.id).unwrap();
1315        old_pending
1316            .response_tx
1317            .send(Ok(RawValue::NULL.to_owned()))
1318            .unwrap();
1319        old_rx.await.unwrap().unwrap();
1320
1321        next_id.store(old_send.id, Ordering::Relaxed);
1322        let (new_tx, new_rx) = tokio::sync::oneshot::channel();
1323        handler.enqueue_request("second", secret_json(json!({})), 0, None, new_tx);
1324        let new_send = send_rx.recv().await.expect("new queued send");
1325        handler.handle_send_failure(&SendFailure {
1326            id: old_send.id,
1327            token: old_send.token,
1328            error: SendError::BrokenPipe("late failure".to_string()),
1329        });
1330
1331        assert_eq!(old_send.id, new_send.id);
1332        assert_ne!(old_send.token, new_send.token);
1333        assert_eq!(
1334            handler.pending.get(&new_send.id).unwrap().token,
1335            new_send.token,
1336        );
1337
1338        handler.handle_send_failure(&SendFailure {
1339            id: new_send.id,
1340            token: new_send.token,
1341            error: SendError::BrokenPipe("current failure".to_string()),
1342        });
1343
1344        let error = new_rx.await.unwrap().expect_err("current request failed");
1345        assert!(error.to_string().contains("current failure"));
1346    }
1347
1348    #[rstest]
1349    #[cfg_attr(all(feature = "simulation", madsim), madsim::test)]
1350    #[cfg_attr(not(all(feature = "simulation", madsim)), tokio::test)]
1351    async fn test_shutdown_aborts_send_worker_and_drains_pending_requests() {
1352        let signal = Arc::new(AtomicBool::new(false));
1353        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1354        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1355        let next_id = Arc::new(AtomicU64::new(1));
1356        let auth_tracker = AuthTracker::new();
1357        let mut handler = feed_handler(
1358            signal,
1359            cmd_rx,
1360            raw_rx,
1361            next_id,
1362            auth_tracker,
1363            unauthenticated_epoch(),
1364        );
1365        let (send_tx, _send_rx) = tokio::sync::mpsc::unbounded_channel();
1366        handler.send_tx = Some(send_tx);
1367
1368        let task = task::spawn(std::future::pending::<()>());
1369        let abort_handle = task.abort_handle();
1370        handler.send_task.insert(task);
1371
1372        let (first_tx, first_rx) = tokio::sync::oneshot::channel();
1373        let (second_tx, second_rx) = tokio::sync::oneshot::channel();
1374        handler.enqueue_request("first", secret_json(json!({})), 0, None, first_tx);
1375        handler.enqueue_request("second", secret_json(json!({})), 0, None, second_tx);
1376        handler.shutdown_send_path("disconnect requested").await;
1377        task::yield_now().await;
1378
1379        let first_error = first_rx.await.unwrap().expect_err("first request failed");
1380        let second_error = second_rx.await.unwrap().expect_err("second request failed");
1381        assert!(abort_handle.is_finished());
1382        assert!(handler.send_tx.is_none());
1383        assert!(handler.send_task.is_empty());
1384        assert!(handler.pending.is_empty());
1385        assert!(first_error.to_string().contains("disconnect requested"));
1386        assert!(second_error.to_string().contains("disconnect requested"));
1387    }
1388
1389    #[rstest]
1390    #[tokio::test]
1391    async fn test_next_stops_when_input_channels_close_with_failure_channel_open() {
1392        let signal = Arc::new(AtomicBool::new(false));
1393        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1394        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1395        let next_id = Arc::new(AtomicU64::new(1));
1396        let auth_tracker = AuthTracker::new();
1397        let mut handler = feed_handler(
1398            signal,
1399            cmd_rx,
1400            raw_rx,
1401            next_id,
1402            auth_tracker,
1403            unauthenticated_epoch(),
1404        );
1405        let (_failure_tx, failure_rx) = tokio::sync::mpsc::unbounded_channel();
1406        handler.send_failure_rx = failure_rx;
1407        drop(cmd_tx);
1408        drop(raw_tx);
1409
1410        let outcome = tokio::time::timeout(std::time::Duration::from_millis(100), handler.next())
1411            .await
1412            .expect("handler stopped after both input channels closed");
1413
1414        assert!(outcome.is_none());
1415    }
1416
1417    #[rstest]
1418    #[tokio::test]
1419    async fn test_uncorrelated_error_fails_only_pending_request() {
1420        let signal = Arc::new(AtomicBool::new(false));
1421        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1422        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1423        let next_id = Arc::new(AtomicU64::new(1));
1424        let auth_tracker = AuthTracker::new();
1425        let mut handler = feed_handler(
1426            signal,
1427            cmd_rx,
1428            raw_rx,
1429            next_id,
1430            auth_tracker,
1431            unauthenticated_epoch(),
1432        );
1433        let (send_tx, _send_rx) = tokio::sync::mpsc::unbounded_channel();
1434        handler.send_tx = Some(send_tx);
1435        let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1436        handler.enqueue_request("request", secret_json(json!({})), 0, None, response_tx);
1437
1438        handler.fail_uncorrelated_error(crate::http::models::JsonRpcError {
1439            code: -32700,
1440            message: "Parse error".to_string(),
1441            data: Some(json!("invalid JSON")),
1442        });
1443
1444        let error = response_rx.await.unwrap().expect_err("request failed");
1445
1446        match error {
1447            DeriveWsError::JsonRpc {
1448                code,
1449                message,
1450                data,
1451            } => {
1452                assert_eq!(code, -32700);
1453                assert_eq!(message, "Parse error");
1454                assert_eq!(data, Some(json!("invalid JSON")));
1455            }
1456            other => panic!("expected JsonRpc, was {other:?}"),
1457        }
1458
1459        assert!(handler.pending.is_empty());
1460    }
1461
1462    #[rstest]
1463    #[tokio::test]
1464    async fn test_uncorrelated_error_does_not_guess_between_pending_requests() {
1465        let signal = Arc::new(AtomicBool::new(false));
1466        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1467        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
1468        let next_id = Arc::new(AtomicU64::new(1));
1469        let auth_tracker = AuthTracker::new();
1470        let mut handler = feed_handler(
1471            signal,
1472            cmd_rx,
1473            raw_rx,
1474            next_id,
1475            auth_tracker,
1476            unauthenticated_epoch(),
1477        );
1478        let (send_tx, _send_rx) = tokio::sync::mpsc::unbounded_channel();
1479        handler.send_tx = Some(send_tx);
1480        let (first_tx, first_rx) = tokio::sync::oneshot::channel();
1481        let (second_tx, second_rx) = tokio::sync::oneshot::channel();
1482        handler.enqueue_request("first", secret_json(json!({})), 0, None, first_tx);
1483        handler.enqueue_request("second", secret_json(json!({})), 0, None, second_tx);
1484
1485        handler.fail_uncorrelated_error(crate::http::models::JsonRpcError {
1486            code: -32600,
1487            message: "Invalid Request".to_string(),
1488            data: None,
1489        });
1490
1491        assert_eq!(handler.pending.len(), 2);
1492        handler.fail_pending("test cleanup");
1493        assert!(first_rx.await.unwrap().is_err());
1494        assert!(second_rx.await.unwrap().is_err());
1495    }
1496}
nautilus_derive/websocket/handler.rs

nautilus_derive/websocket/
handler.rs
