1use std::{
25 fmt::Debug,
26 sync::{
27 Arc,
28 atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
29 },
30 time::Duration,
31};
32
33use alloy::signers::local::PrivateKeySigner;
34use arc_swap::ArcSwap;
35use nautilus_common::live::dst::time;
36#[cfg(test)]
37use nautilus_common::live::get_runtime;
38use nautilus_core::{
39 UUID4,
40 string::secret::{REDACTED, SecretString},
41};
42use nautilus_live::{
43 SocketControl,
44 book::snapshot::SnapshotGate,
45 task::{SharedTaskSlot, TaskJoinOutcome, TaskSlot, finish_task},
46};
47use nautilus_network::{
48 ConnectionUrl,
49 http::create_standard_nautilus_headers,
50 mode::ConnectionMode,
51 ratelimiter::clock::MonotonicClock,
52 websocket::{
53 AuthTracker, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
54 channel_message_handler,
55 },
56};
57use serde::{Deserialize, Serialize, de::DeserializeOwned};
58use serde_json::value::RawValue;
59use tokio_util::sync::CancellationToken;
60use ustr::Ustr;
61use zeroize::Zeroizing;
62
63use super::{
64 error::{DeriveWsError, Result},
65 handler::{
66 DeriveWsMessage, FeedHandler, HandlerCommand, orderbook_subscribe_params,
67 ticker_subscribe_params, trades_subscribe_params,
68 },
69 messages::{
70 DeriveWsChannel, WsLoginParams, WsLoginResult, WsSubscribeParams, WsSubscribeResult,
71 WsUnsubscribeParams, WsUnsubscribeResult, methods, orderbook_channel, ticker_channel,
72 trades_channel,
73 },
74};
75use crate::{
76 book::{recovery::BookRecovery, sync::BookSyncTracker},
77 common::{
78 consts::{
79 RECONNECT_BACKOFF_FACTOR, RECONNECT_BASE_BACKOFF, RECONNECT_JITTER_MS,
80 RECONNECT_MAX_BACKOFF, RECONNECT_TIMEOUT, WS_HEARTBEAT_SECS, WS_HEARTBEAT_TIMEOUT,
81 WS_REQUEST_TIMEOUT,
82 },
83 enums::DeriveEnvironment,
84 parse::salvage_rows,
85 rate_limit::{DeriveRateLimiter, RateClass, RateReservation, rate_class_for_method},
86 retry::should_retry_ws_error,
87 urls,
88 },
89 http::{
90 models::{
91 DeriveCancelByInstrumentResult, DeriveCancelByLabelResult, DeriveEmptyResult,
92 DeriveOpenOrdersResult, DeriveOrder, DeriveOrderResult, DeriveReplaceOutcome,
93 DeriveReplaceResult, DeriveTrade,
94 },
95 query::{
96 DeriveCancelAllParams, DeriveCancelByInstrumentParams, DeriveCancelByLabelParams,
97 DeriveCancelParams, DeriveCancelTriggerOrderParams, DeriveGetTriggerOrdersParams,
98 DeriveOrderParams, DeriveReplaceParams, DeriveTriggerOrderParams,
99 },
100 },
101 signing::auth::build_ws_login,
102};
103
104#[derive(Clone)]
107pub struct DeriveWsCredentials {
108 pub wallet_address: String,
110 pub signer: PrivateKeySigner,
112}
113
114impl DeriveWsCredentials {
115 pub fn new(wallet_address: impl Into<String>, session_key_hex: &str) -> Result<Self> {
121 let signer: PrivateKeySigner = session_key_hex
122 .parse()
123 .map_err(|e| DeriveWsError::transport(format!("invalid session key: {e}")))?;
124 Ok(Self {
125 wallet_address: wallet_address.into(),
126 signer,
127 })
128 }
129}
130
131impl Debug for DeriveWsCredentials {
132 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
133 f.debug_struct(stringify!(DeriveWsCredentials))
134 .field("wallet_address", &self.wallet_address)
135 .field("signer", &REDACTED)
136 .finish()
137 }
138}
139
140type WsRateLimiter = DeriveRateLimiter;
143
144const MAX_SESSION_RECOVERY_ATTEMPTS: u32 = 3;
145const SUBSCRIPTION_ACCEPTED_STATUSES: &[&str] = &["ok"];
146const SUBSCRIPTION_REPLAY_ACCEPTED_STATUSES: &[&str] = &["ok", "already subscribed"];
147pub(super) const UNAUTHENTICATED_CONNECTION_EPOCH: u64 = u64::MAX;
148
149#[derive(Debug)]
155pub struct DeriveWebSocketClient {
156 url: ConnectionUrl,
157 transport_backend: TransportBackend,
158 proxy_url: Option<SecretString>,
159 connection_mode: Arc<ArcSwap<AtomicU8>>,
160 connection_epoch: Arc<ArcSwap<AtomicU64>>,
161 signal: Arc<AtomicBool>,
162 auth_tracker: AuthTracker,
163 authenticated_epoch: Arc<AtomicU64>,
164 credentials: Option<DeriveWsCredentials>,
165 next_id: Arc<AtomicU64>,
166 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
167 out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<DeriveWsMessage>>,
168 subscriptions: Arc<SubscriptionState>,
169 subscription_lock: Arc<tokio::sync::Mutex<()>>,
170 book_sync: Option<(BookSyncTracker, Duration)>,
171 task_handle: TaskSlot<()>,
172 send_task: Arc<SharedTaskSlot<()>>,
173 recovery_task: Arc<SharedTaskSlot<()>>,
174 shutdown_errors: Vec<String>,
175 request_timeout: Duration,
176 conn_id: Arc<ArcSwap<String>>,
177 rate_limiter: Arc<WsRateLimiter>,
178 socket_control: Option<SocketControl>,
179}
180
181#[derive(Debug, Clone)]
183pub struct DeriveWebSocketSubscriptionHandle {
184 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
185 subscriptions: Arc<SubscriptionState>,
186 subscription_lock: Arc<tokio::sync::Mutex<()>>,
187 book_sync: Option<(BookSyncTracker, Duration)>,
188 connection_epoch: Arc<ArcSwap<AtomicU64>>,
189 request_timeout: Duration,
190 rate_limiter: Arc<WsRateLimiter>,
191}
192
193#[derive(Debug, Clone)]
202pub struct DeriveWsExecutionHandle {
203 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
204 auth_tracker: AuthTracker,
205 request_timeout: Duration,
206 conn_id: Arc<ArcSwap<String>>,
207 rate_limiter: Arc<WsRateLimiter>,
208}
209
210#[derive(Debug)]
211pub(crate) struct MatchingRateLimitReservation {
212 method: &'static str,
213 quota: RateReservation<MonotonicClock>,
214}
215
216impl DeriveWebSocketClient {
217 #[must_use]
220 pub fn new(
221 url: Option<String>,
222 environment: DeriveEnvironment,
223 transport_backend: TransportBackend,
224 proxy_url: Option<String>,
225 ) -> Self {
226 let url = url.unwrap_or_else(|| urls::ws_url(environment).to_string());
227 Self::build(url, transport_backend, proxy_url, None, None, None)
228 }
229
230 #[must_use]
241 pub fn with_credentials(
242 url: Option<String>,
243 environment: DeriveEnvironment,
244 transport_backend: TransportBackend,
245 proxy_url: Option<String>,
246 credentials: DeriveWsCredentials,
247 max_matching_requests_per_second: Option<u32>,
248 max_per_instrument_matching_requests_per_second: Option<u32>,
249 ) -> Self {
250 let url = url.unwrap_or_else(|| urls::ws_url(environment).to_string());
251 Self::build(
252 url,
253 transport_backend,
254 proxy_url,
255 Some(credentials),
256 max_matching_requests_per_second,
257 max_per_instrument_matching_requests_per_second,
258 )
259 }
260
261 fn build(
262 url: String,
263 transport_backend: TransportBackend,
264 proxy_url: Option<String>,
265 credentials: Option<DeriveWsCredentials>,
266 matching_tps: Option<u32>,
267 instrument_tps: Option<u32>,
268 ) -> Self {
269 let connection_mode = Arc::new(ArcSwap::new(Arc::new(AtomicU8::new(
270 ConnectionMode::Closed as u8,
271 ))));
272 let connection_epoch = Arc::new(ArcSwap::new(Arc::new(AtomicU64::new(0))));
273
274 let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
276
277 let rate_limiter = Arc::new(DeriveRateLimiter::shared(
278 &url,
279 credentials
280 .as_ref()
281 .map(|credentials| credentials.wallet_address.as_str()),
282 matching_tps,
283 instrument_tps,
284 ));
285
286 Self {
287 url: url.into(),
288 transport_backend,
289 proxy_url: proxy_url.map(SecretString::from),
290 connection_mode,
291 connection_epoch,
292 signal: Arc::new(AtomicBool::new(false)),
293 auth_tracker: AuthTracker::new(),
294 authenticated_epoch: Arc::new(AtomicU64::new(UNAUTHENTICATED_CONNECTION_EPOCH)),
295 credentials,
296 next_id: Arc::new(AtomicU64::new(1)),
297 cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
298 out_rx: None,
299 subscriptions: Arc::new(SubscriptionState::new('.')),
300 subscription_lock: Arc::new(tokio::sync::Mutex::new(())),
301 book_sync: None,
302 task_handle: TaskSlot::new(),
303 send_task: Arc::new(SharedTaskSlot::new()),
304 recovery_task: Arc::new(SharedTaskSlot::new()),
305 shutdown_errors: Vec::new(),
306 request_timeout: WS_REQUEST_TIMEOUT,
307 conn_id: Arc::new(ArcSwap::from_pointee(UUID4::new().to_string())),
308 rate_limiter,
309 socket_control: None,
310 }
311 }
312
313 #[must_use]
314 pub(crate) fn with_book_sync(mut self, tracker: BookSyncTracker, timeout: Duration) -> Self {
315 self.book_sync = Some((tracker, timeout));
316 self
317 }
318
319 #[must_use]
321 pub fn with_socket_control(mut self, control: SocketControl) -> Self {
322 self.socket_control = Some(control);
323 self
324 }
325
326 #[must_use]
328 pub fn url(&self) -> &str {
329 self.url.expose()
330 }
331
332 pub fn set_request_timeout(&mut self, timeout: Duration) {
336 self.request_timeout = timeout;
337 }
338
339 #[must_use]
342 pub fn is_authenticated(&self) -> bool {
343 self.auth_tracker.is_authenticated()
344 && self.is_active()
345 && self.authenticated_epoch.load(Ordering::Acquire)
346 == self.connection_epoch.load().load(Ordering::Acquire)
347 }
348
349 #[must_use]
351 pub fn is_active(&self) -> bool {
352 self.connection_mode.load().load(Ordering::Relaxed) == ConnectionMode::Active as u8
353 }
354
355 pub async fn connect(&mut self) -> Result<()> {
364 let auth_ok = self.credentials.is_none() || self.is_authenticated();
367 if self.is_active() && auth_ok && self.task_handle.is_some() {
368 log::warn!("Derive WebSocket already connected");
369 return Ok(());
370 }
371
372 if self.task_handle.is_some()
374 || !self.send_task.is_empty()
375 || !self.recovery_task.is_empty()
376 {
377 log::debug!("Tearing down stale Derive WebSocket state before connect");
378 self.teardown().await?;
379 }
380
381 self.signal.store(false, Ordering::Release);
382 let setup_guard = DeriveWebSocketSetupGuard::new(self.shutdown_handle());
383
384 self.authenticated_epoch
385 .store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
386
387 let (message_handler, raw_rx) = channel_message_handler();
388 let headers = create_standard_nautilus_headers();
389
390 let cfg = WebSocketConfig {
391 url: self.url.expose().to_owned(),
392 headers,
393 heartbeat_interval_secs: Some(WS_HEARTBEAT_SECS),
394 heartbeat_payload: None,
395 connect_timeout_ms: Some(RECONNECT_TIMEOUT.as_millis() as u64),
396 reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
397 reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
398 reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
399 reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
400 reconnect_max_attempts: None,
401 heartbeat_timeout_secs: Some(WS_HEARTBEAT_TIMEOUT.as_secs()),
402 idle_timeout_ms: None,
403 writer_capacity: None,
404 backend: self.transport_backend,
405 proxy_url: self
406 .proxy_url
407 .as_ref()
408 .map(|value| value.expose_secret().to_owned()),
409 max_message_size_bytes: None,
410 max_frame_size_bytes: None,
411 };
412
413 let client = WebSocketClient::builder()
417 .config(cfg)
418 .message_handler(message_handler)
419 .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
420 .connect()
421 .await
422 .map_err(|e| DeriveWsError::transport(e.to_string()))?;
423
424 client.set_auth_tracker(self.auth_tracker.clone(), false);
428
429 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
430 let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<DeriveWsMessage>();
431
432 *self.cmd_tx.write().await = cmd_tx.clone();
433 self.out_rx = Some(out_rx);
434 self.conn_id.store(Arc::new(UUID4::new().to_string()));
435
436 let connection_mode = client.connection_mode_atomic();
437 let connection_epoch = client.connection_epoch_atomic();
438 let reconnect_handle = client.reconnect_handle();
439 self.connection_mode.store(Arc::clone(&connection_mode));
440 self.connection_epoch.store(Arc::clone(&connection_epoch));
441 log::debug!("Derive WebSocket connected: {}", self.url.origin());
442
443 if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
444 return Err(DeriveWsError::transport(format!(
445 "failed to send SetClient command: {e}",
446 )));
447 }
448
449 let recovery = SessionRecovery {
450 subscription: self.subscription_handle(),
451 cmd_tx: cmd_tx.clone(),
452 auth_tracker: self.auth_tracker.clone(),
453 authenticated_epoch: Arc::clone(&self.authenticated_epoch),
454 connection_mode: Arc::clone(&connection_mode),
455 connection_epoch: Arc::clone(&connection_epoch),
456 credentials: self.credentials.clone(),
457 signal: Arc::clone(&self.signal),
458 };
459
460 let handler = FeedHandler::new_with_send_task(
461 Arc::clone(&self.signal),
462 cmd_rx,
463 raw_rx,
464 Arc::clone(&self.next_id),
465 self.auth_tracker.clone(),
466 Arc::clone(&self.authenticated_epoch),
467 Arc::clone(&self.send_task),
468 )
469 .with_book_sync(self.book_sync.clone());
470
471 if let Err(e) = self.task_handle.spawn(run_session_handler(
472 handler,
473 recovery,
474 out_tx,
475 Arc::clone(&self.recovery_task),
476 Arc::clone(&self.conn_id),
477 )) {
478 let shutdown_result = self.teardown().await;
479 return Err(DeriveWsError::transport(match shutdown_result {
480 Ok(()) => format!("failed to start WebSocket handler task: {e}"),
481 Err(shutdown_error) => format!(
482 "failed to start WebSocket handler task: {e}; startup rollback failed: \
483 {shutdown_error}"
484 ),
485 }));
486 }
487
488 if let Some(control) = &self.socket_control {
489 control.register(move || reconnect_handle.request_reconnect());
490 }
491
492 if let Some(creds) = self.credentials.clone()
493 && let Err(e) = login_via_handler(
494 &self.rate_limiter,
495 &cmd_tx,
496 &self.auth_tracker,
497 &self.authenticated_epoch,
498 &connection_mode,
499 &connection_epoch,
500 &creds,
501 self.request_timeout,
502 )
503 .await
504 {
505 log::warn!(
508 "Derive WebSocket login failed; tearing down transport: {}",
509 e.redacted_message()
510 );
511 self.teardown().await?;
512 return Err(e);
513 }
514
515 setup_guard.disarm();
516 Ok(())
517 }
518
519 pub(crate) fn begin_shutdown(&self) {
520 self.shutdown_handle().begin_shutdown();
521 }
522
523 pub(crate) fn shutdown_handle(&self) -> DeriveWebSocketShutdownHandle {
524 DeriveWebSocketShutdownHandle {
525 signal: Arc::clone(&self.signal),
526 recovery_task: Arc::clone(&self.recovery_task),
527 }
528 }
529
530 async fn teardown(&mut self) -> Result<()> {
534 self.begin_shutdown();
535
536 if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
537 log::debug!(
538 "Failed to enqueue Disconnect command (handler may already be shut down): {e}",
539 );
540 }
541
542 match finish_task(
543 &mut self.task_handle,
544 Duration::from_secs(2),
545 Duration::from_secs(2),
546 )
547 .await
548 {
549 None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => {}
550 Some(TaskJoinOutcome::Failed(error)) => self
551 .shutdown_errors
552 .push(format!("WebSocket handler task failed: {error}")),
553 Some(TaskJoinOutcome::Incomplete) => self
554 .shutdown_errors
555 .push("WebSocket handler task did not stop after abort".to_string()),
556 }
557
558 if self.task_handle.is_some() {
559 if let Some(control) = &self.socket_control {
560 control.deregister();
561 }
562
563 return self.take_shutdown_result();
564 }
565
566 for (name, slot) in [
567 ("recovery", Arc::clone(&self.recovery_task)),
568 ("send", Arc::clone(&self.send_task)),
569 ] {
570 match slot
571 .finish(Duration::from_secs(1), Duration::from_secs(2))
572 .await
573 {
574 None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => {}
575 Some(TaskJoinOutcome::Failed(error)) => self
576 .shutdown_errors
577 .push(format!("WebSocket {name} worker failed: {error}")),
578 Some(TaskJoinOutcome::Incomplete) => self
579 .shutdown_errors
580 .push(format!("WebSocket {name} worker did not stop after abort")),
581 }
582
583 if !slot.is_empty() {
584 if let Some(control) = &self.socket_control {
585 control.deregister();
586 }
587
588 return self.take_shutdown_result();
589 }
590 }
591
592 let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
595 *self.cmd_tx.write().await = placeholder_tx;
596 self.out_rx = None;
597 self.connection_mode
598 .store(Arc::new(AtomicU8::new(ConnectionMode::Closed as u8)));
599 self.connection_epoch.store(Arc::new(AtomicU64::new(0)));
600 self.auth_tracker.invalidate();
601 self.authenticated_epoch
602 .store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
603 self.subscriptions.clear();
604
605 if let Some((tracker, _)) = &self.book_sync {
606 tracker.clear();
607 }
608
609 self.signal.store(false, Ordering::Relaxed);
610
611 if let Some(control) = &self.socket_control {
612 control.deregister();
613 }
614
615 self.take_shutdown_result()
616 }
617
618 fn take_shutdown_result(&mut self) -> Result<()> {
619 if self.shutdown_errors.is_empty() {
620 Ok(())
621 } else {
622 Err(DeriveWsError::transport(
623 std::mem::take(&mut self.shutdown_errors).join("; "),
624 ))
625 }
626 }
627
628 pub async fn disconnect(&mut self) -> Result<()> {
635 log::debug!("Disconnecting Derive WebSocket");
636 self.begin_shutdown();
637 self.teardown().await
638 }
639
640 pub async fn subscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
647 self.subscription_handle()
648 .subscribe_ticker(instrument_name, interval)
649 .await
650 }
651
652 pub async fn unsubscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
658 self.subscription_handle()
659 .unsubscribe_ticker(instrument_name, interval)
660 .await
661 }
662
663 pub async fn subscribe_orderbook(
669 &self,
670 instrument_name: &str,
671 group: &str,
672 depth: &str,
673 ) -> Result<()> {
674 self.subscription_handle()
675 .subscribe_orderbook(instrument_name, group, depth)
676 .await
677 }
678
679 pub async fn unsubscribe_orderbook(
685 &self,
686 instrument_name: &str,
687 group: &str,
688 depth: &str,
689 ) -> Result<()> {
690 self.subscription_handle()
691 .unsubscribe_orderbook(instrument_name, group, depth)
692 .await
693 }
694
695 pub async fn subscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
701 self.subscription_handle()
702 .subscribe_trades(instrument_type, currency)
703 .await
704 }
705
706 pub async fn unsubscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
712 self.subscription_handle()
713 .unsubscribe_trades(instrument_type, currency)
714 .await
715 }
716
717 pub async fn subscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
727 where
728 C: Into<DeriveWsChannel>,
729 {
730 self.subscription_handle()
731 .subscribe_channels(channels)
732 .await
733 }
734
735 pub async fn unsubscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
742 where
743 C: Into<DeriveWsChannel>,
744 {
745 self.subscription_handle()
746 .unsubscribe_channels(channels)
747 .await
748 }
749
750 pub async fn next_event(&mut self) -> Option<DeriveWsMessage> {
752 if let Some(rx) = self.out_rx.as_mut() {
753 rx.recv().await
754 } else {
755 None
756 }
757 }
758
759 #[must_use]
762 pub fn subscription_count(&self) -> usize {
763 self.subscriptions.len()
764 }
765
766 #[must_use]
768 pub fn subscription_handle(&self) -> DeriveWebSocketSubscriptionHandle {
769 DeriveWebSocketSubscriptionHandle {
770 cmd_tx: Arc::clone(&self.cmd_tx),
771 subscriptions: Arc::clone(&self.subscriptions),
772 subscription_lock: Arc::clone(&self.subscription_lock),
773 book_sync: self.book_sync.clone(),
774 connection_epoch: Arc::clone(&self.connection_epoch),
775 request_timeout: self.request_timeout,
776 rate_limiter: Arc::clone(&self.rate_limiter),
777 }
778 }
779
780 #[must_use]
787 pub fn execution_handle(&self) -> DeriveWsExecutionHandle {
788 DeriveWsExecutionHandle {
789 cmd_tx: Arc::clone(&self.cmd_tx),
790 auth_tracker: self.auth_tracker.clone(),
791 request_timeout: self.request_timeout,
792 conn_id: Arc::clone(&self.conn_id),
793 rate_limiter: Arc::clone(&self.rate_limiter),
794 }
795 }
796
797 pub fn take_event_receiver(
802 &mut self,
803 ) -> Option<tokio::sync::mpsc::UnboundedReceiver<DeriveWsMessage>> {
804 self.out_rx.take()
805 }
806}
807
808impl Drop for DeriveWebSocketClient {
809 fn drop(&mut self) {
810 self.begin_shutdown();
811
812 if let Some(handle) = self.task_handle.as_ref() {
813 handle.abort();
814 }
815
816 if let Some(control) = &self.socket_control {
817 control.deregister();
818 }
819 }
820}
821
822struct SessionRecovery {
823 subscription: DeriveWebSocketSubscriptionHandle,
824 cmd_tx: tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
825 auth_tracker: AuthTracker,
826 authenticated_epoch: Arc<AtomicU64>,
827 connection_mode: Arc<AtomicU8>,
828 connection_epoch: Arc<AtomicU64>,
829 credentials: Option<DeriveWsCredentials>,
830 signal: Arc<AtomicBool>,
831}
832
833async fn run_session_handler(
834 mut handler: FeedHandler,
835 recovery: SessionRecovery,
836 out_tx: tokio::sync::mpsc::UnboundedSender<DeriveWsMessage>,
837 recovery_task: Arc<SharedTaskSlot<()>>,
838 conn_id: Arc<ArcSwap<String>>,
839) {
840 let (recovery_tx, recovery_rx) = tokio::sync::mpsc::unbounded_channel();
841 let cmd_tx = recovery.cmd_tx.clone();
842 let connection_epoch = Arc::clone(&recovery.connection_epoch);
843 if let Err(e) = recovery_task.spawn(recovery.run(recovery_rx, out_tx.clone())) {
844 let _ = out_tx.send(DeriveWsMessage::SessionRecoveryFailed(format!(
845 "failed to start recovery worker: {e}"
846 )));
847 let _ = cmd_tx.send(HandlerCommand::Disconnect);
848 return;
849 }
850
851 loop {
852 match handler.next().await {
853 Some(DeriveWsMessage::Reconnected) => {
854 log::info!("Derive WebSocket re-establishing session after reconnect");
855 conn_id.store(Arc::new(UUID4::new().to_string()));
856 let epoch = connection_epoch.load(Ordering::Acquire);
857 if recovery_tx.send(epoch).is_err() {
858 log::error!("Derive WebSocket recovery task stopped unexpectedly");
859 let _ = cmd_tx.send(HandlerCommand::Disconnect);
860 }
861 }
862 Some(msg) => {
863 if out_tx.send(msg).is_err() {
864 log::debug!("Derive outer receiver dropped, exiting stream loop");
865 break;
866 }
867 }
868 None => {
869 log::debug!("Derive handler task ended");
870 break;
871 }
872 }
873 }
874
875 recovery_task.abort();
876}
877
878impl SessionRecovery {
879 async fn run(
880 self,
881 mut rx: tokio::sync::mpsc::UnboundedReceiver<u64>,
882 out_tx: tokio::sync::mpsc::UnboundedSender<DeriveWsMessage>,
883 ) {
884 if self.signal.load(Ordering::Acquire) {
885 return;
886 }
887
888 while let Some(mut requested_epoch) = rx.recv().await {
889 take_latest_epoch(&mut rx, &mut requested_epoch);
890 match self.recover_until_current(&mut rx, requested_epoch).await {
891 Ok(()) => {
892 if out_tx.send(DeriveWsMessage::Reconnected).is_err() {
893 log::debug!("Derive outer receiver dropped during recovery");
894 }
895 }
896 Err(e) => {
897 log::error!(
898 "Derive WebSocket session recovery failed: {}",
899 e.redacted_message()
900 );
901 let _ = out_tx.send(DeriveWsMessage::SessionRecoveryFailed(e.to_string()));
902 let _ = self.cmd_tx.send(HandlerCommand::Disconnect);
903 return;
904 }
905 }
906 }
907 }
908
909 async fn recover_until_current(
910 &self,
911 rx: &mut tokio::sync::mpsc::UnboundedReceiver<u64>,
912 mut requested_epoch: u64,
913 ) -> Result<()> {
914 loop {
915 match recover_session(
916 &self.subscription.rate_limiter,
917 &self.cmd_tx,
918 &self.auth_tracker,
919 &self.authenticated_epoch,
920 &self.connection_mode,
921 &self.connection_epoch,
922 self.credentials.as_ref(),
923 &self.subscription.subscriptions,
924 &self.subscription.subscription_lock,
925 self.subscription.request_timeout,
926 )
927 .await
928 {
929 Ok(recovered_epoch) => {
930 take_latest_epoch(rx, &mut requested_epoch);
931 if requested_epoch > recovered_epoch
932 || !session_connection_is_active(
933 &self.connection_mode,
934 &self.connection_epoch,
935 recovered_epoch,
936 )
937 {
938 continue;
939 }
940
941 return Ok(());
942 }
943 Err(e) => {
944 if take_latest_epoch(rx, &mut requested_epoch) {
945 continue;
946 }
947
948 return Err(e);
949 }
950 }
951 }
952 }
953}
954
955fn take_latest_epoch(
956 rx: &mut tokio::sync::mpsc::UnboundedReceiver<u64>,
957 requested_epoch: &mut u64,
958) -> bool {
959 let mut received = false;
960
961 while let Ok(epoch) = rx.try_recv() {
962 *requested_epoch = (*requested_epoch).max(epoch);
963 received = true;
964 }
965
966 received
967}
968
969#[derive(Clone, Debug)]
970pub(crate) struct DeriveWebSocketShutdownHandle {
971 signal: Arc<AtomicBool>,
972 recovery_task: Arc<SharedTaskSlot<()>>,
973}
974
975impl DeriveWebSocketShutdownHandle {
976 pub(crate) fn begin_shutdown(&self) {
977 self.signal.store(true, Ordering::Release);
978 self.recovery_task.abort();
979 }
980}
981
982struct DeriveWebSocketSetupGuard {
983 shutdown: DeriveWebSocketShutdownHandle,
984 armed: bool,
985}
986
987impl DeriveWebSocketSetupGuard {
988 fn new(shutdown: DeriveWebSocketShutdownHandle) -> Self {
989 Self {
990 shutdown,
991 armed: true,
992 }
993 }
994
995 fn disarm(mut self) {
996 self.armed = false;
997 }
998}
999
1000impl Drop for DeriveWebSocketSetupGuard {
1001 fn drop(&mut self) {
1002 if self.armed {
1003 self.shutdown.begin_shutdown();
1004 }
1005 }
1006}
1007
1008impl DeriveWebSocketSubscriptionHandle {
1009 pub(crate) async fn subscribe_book_gated(&self, channel: String) -> Result<()> {
1010 let _guard = self.subscription_lock.lock().await;
1011 self.book_request_locked(methods::PUBLIC_SUBSCRIBE, channel, None, None)
1012 .await
1013 }
1014
1015 pub(crate) async fn replace_book_gated(
1016 &self,
1017 channel: String,
1018 cancel: CancellationToken,
1019 recovery: &Arc<BookRecovery>,
1020 gate: SnapshotGate,
1021 ) -> Result<()> {
1022 let _guard = tokio::select! {
1023 biased;
1024 () = cancel.cancelled() => return Err(DeriveWsError::NotConnected),
1025 guard = self.subscription_lock.lock() => guard,
1026 };
1027 self.book_sync
1028 .as_ref()
1029 .and_then(|(tracker, _)| tracker.begin_recovery_write(&channel, recovery, gate))
1030 .ok_or(DeriveWsError::NotConnected)?;
1031 let epoch = self.connection_epoch.load().load(Ordering::Acquire);
1032
1033 self.remember_subscription(&channel);
1035 self.book_request_locked(
1036 methods::PUBLIC_UNSUBSCRIBE,
1037 channel.clone(),
1038 Some(cancel.clone()),
1039 Some(epoch),
1040 )
1041 .await?;
1042 self.book_request_locked(
1043 methods::PUBLIC_SUBSCRIBE,
1044 channel,
1045 Some(cancel),
1046 Some(epoch),
1047 )
1048 .await
1049 }
1050
1051 async fn book_request_locked(
1052 &self,
1053 method: &'static str,
1054 channel: String,
1055 cancel: Option<CancellationToken>,
1056 epoch: Option<u64>,
1057 ) -> Result<()> {
1058 let params = WsSubscribeParams {
1059 channels: vec![DeriveWsChannel::from(channel.clone())],
1060 };
1061
1062 self.remember_subscription(&channel);
1063 self.rate_limiter
1064 .await_class_ready(rate_class_for_method(method), None)
1065 .await;
1066
1067 let write = self
1068 .book_sync
1069 .as_ref()
1070 .and_then(|(tracker, _)| tracker.write(&channel))
1071 .ok_or(DeriveWsError::NotConnected)?;
1072 let generation = write.generation.clone();
1073 let cancel = cancel.unwrap_or_else(|| generation.clone());
1074 let epoch = epoch.unwrap_or_else(|| self.connection_epoch.load().load(Ordering::Acquire));
1075
1076 if cancel.is_cancelled() || generation.is_cancelled() {
1077 return Err(DeriveWsError::NotConnected);
1078 }
1079
1080 let cmd_tx = self.cmd_tx.read().await.clone();
1081 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1082 cmd_tx
1083 .send(HandlerCommand::Request {
1084 method,
1085 params: SecretString::from(serde_json::to_string(¶ms)?),
1086 connection_epoch: Some(epoch),
1087 cancellation: Some(cancel.clone()),
1088 response_tx,
1089 })
1090 .map_err(|_| DeriveWsError::NotConnected)?;
1091
1092 let result = tokio::select! {
1093 biased;
1094 () = cancel.cancelled() => return Err(DeriveWsError::NotConnected),
1095 () = generation.cancelled() => return Err(DeriveWsError::NotConnected),
1096 result = time::timeout(self.request_timeout, response_rx) => result,
1097 };
1098
1099 let value = result
1100 .map_err(|_| DeriveWsError::Timeout {
1101 method: method.to_owned(),
1102 })?
1103 .map_err(|_| DeriveWsError::RequestCancelled {
1104 method: method.to_owned(),
1105 })??;
1106
1107 if method == methods::PUBLIC_SUBSCRIBE {
1108 let result: WsSubscribeResult = decode_default_result(&value)?;
1109 subscription_outcome(
1110 ¶ms.channels,
1111 &result,
1112 SUBSCRIPTION_REPLAY_ACCEPTED_STATUSES,
1113 &self.subscriptions,
1114 )
1115 } else {
1116 let result: WsUnsubscribeResult = decode_default_result(&value)?;
1117
1118 let topic = DeriveWsChannel::from(channel.clone());
1120
1121 if result
1122 .status
1123 .get(&topic)
1124 .is_some_and(|status| status.as_str() == "ok")
1125 && !result.remaining_subscriptions.contains(&topic)
1126 {
1127 Ok(())
1128 } else {
1129 Err(DeriveWsError::Subscription {
1130 details: format!("{channel}: unsubscribe not confirmed"),
1131 })
1132 }
1133 }
1134 }
1135
1136 pub(crate) fn has_subscription(&self, channel: &str) -> bool {
1137 let (name, symbol) = channel.split_once('.').unwrap_or((channel, ""));
1138 self.subscriptions
1139 .is_subscribed(&Ustr::from(name), &Ustr::from(symbol))
1140 }
1141
1142 pub(crate) fn forget_subscription(&self, channel: &str) {
1143 self.subscriptions.mark_unsubscribe(channel);
1144 self.subscriptions.confirm_unsubscribe(channel);
1145 }
1146
1147 pub(crate) fn remember_subscription(&self, channel: &str) {
1148 self.subscriptions.mark_subscribe(channel);
1149 }
1150
1151 pub async fn subscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
1157 let params = ticker_subscribe_params(instrument_name, interval);
1158 self.subscribe_channels(params.channels).await
1159 }
1160
1161 pub async fn unsubscribe_ticker(&self, instrument_name: &str, interval: &str) -> Result<()> {
1167 let channel = ticker_channel(instrument_name, interval);
1168 self.send_unsubscribe(channel).await
1169 }
1170
1171 pub async fn subscribe_orderbook(
1177 &self,
1178 instrument_name: &str,
1179 group: &str,
1180 depth: &str,
1181 ) -> Result<()> {
1182 let params = orderbook_subscribe_params(instrument_name, group, depth);
1183 self.subscribe_channels(params.channels).await
1184 }
1185
1186 pub async fn unsubscribe_orderbook(
1192 &self,
1193 instrument_name: &str,
1194 group: &str,
1195 depth: &str,
1196 ) -> Result<()> {
1197 let channel = orderbook_channel(instrument_name, group, depth);
1198 self.send_unsubscribe(channel).await
1199 }
1200
1201 pub async fn subscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
1207 let params = trades_subscribe_params(instrument_type, currency);
1208 self.subscribe_channels(params.channels).await
1209 }
1210
1211 pub async fn unsubscribe_trades(&self, instrument_type: &str, currency: &str) -> Result<()> {
1217 let channel = trades_channel(instrument_type, currency);
1218 self.send_unsubscribe(channel).await
1219 }
1220
1221 pub async fn subscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
1227 where
1228 C: Into<DeriveWsChannel>,
1229 {
1230 let channels = channels.into_iter().map(Into::into).collect::<Vec<_>>();
1231 if channels.is_empty() {
1232 return Ok(());
1233 }
1234
1235 let _guard = self.subscription_lock.lock().await;
1236 let params = WsSubscribeParams { channels };
1237 for topic in channel_topics(¶ms.channels) {
1238 self.subscriptions.mark_subscribe(&topic);
1239 }
1240
1241 let cmd_tx = self.cmd_tx.read().await.clone();
1242 let result: WsSubscribeResult = send_request(
1243 &self.rate_limiter,
1244 &cmd_tx,
1245 methods::PUBLIC_SUBSCRIBE,
1246 ¶ms,
1247 self.request_timeout,
1248 )
1249 .await?;
1250
1251 subscription_outcome(
1252 ¶ms.channels,
1253 &result,
1254 SUBSCRIPTION_ACCEPTED_STATUSES,
1255 &self.subscriptions,
1256 )
1257 }
1258
1259 pub async fn unsubscribe_channels<C>(&self, channels: Vec<C>) -> Result<()>
1266 where
1267 C: Into<DeriveWsChannel>,
1268 {
1269 let channels = channels.into_iter().map(Into::into).collect::<Vec<_>>();
1270 if channels.is_empty() {
1271 return Ok(());
1272 }
1273
1274 let _guard = self.subscription_lock.lock().await;
1275
1276 for topic in channel_topics(&channels) {
1277 self.subscriptions.mark_unsubscribe(&topic);
1278 }
1279
1280 let params = WsUnsubscribeParams { channels };
1281 let cmd_tx = self.cmd_tx.read().await.clone();
1282 let result: WsUnsubscribeResult = send_request(
1283 &self.rate_limiter,
1284 &cmd_tx,
1285 methods::PUBLIC_UNSUBSCRIBE,
1286 ¶ms,
1287 self.request_timeout,
1288 )
1289 .await?;
1290
1291 unsubscribe_outcome(¶ms.channels, &result, &self.subscriptions)
1292 }
1293
1294 async fn send_unsubscribe(&self, channel: String) -> Result<()> {
1295 self.unsubscribe_channels(vec![DeriveWsChannel::from(channel)])
1296 .await
1297 }
1298}
1299
1300impl DeriveWsExecutionHandle {
1301 #[must_use]
1303 pub fn conn_id(&self) -> String {
1304 self.conn_id.load_full().as_ref().clone()
1305 }
1306
1307 pub async fn submit_order(&self, params: &DeriveOrderParams) -> Result<DeriveOrder> {
1318 let reservation = self
1319 .reserve_matching_request(methods::PRIVATE_ORDER, ¶ms.instrument_name)
1320 .await?;
1321 self.submit_order_after_rate_limit(params, reservation)
1322 .await
1323 .map(|result| result.order)
1324 }
1325
1326 pub(crate) async fn submit_order_after_rate_limit(
1327 &self,
1328 params: &DeriveOrderParams,
1329 reservation: MatchingRateLimitReservation,
1330 ) -> Result<DeriveOrderResult> {
1331 self.ensure_authenticated(methods::PRIVATE_ORDER)?;
1332 debug_assert_eq!(reservation.method, methods::PRIVATE_ORDER);
1333 let cmd_tx = self.cmd_tx.read().await.clone();
1334 self.ensure_authenticated(reservation.method)?;
1335 self.commit_matching_reservation(reservation)?;
1336 let result: DeriveOrderResult = send_request_typed_after_rate_limit(
1337 &self.rate_limiter,
1338 &cmd_tx,
1339 methods::PRIVATE_ORDER,
1340 params,
1341 self.request_timeout,
1342 )
1343 .await?;
1344 Ok(result)
1345 }
1346
1347 pub async fn submit_trigger_order(
1354 &self,
1355 params: &DeriveTriggerOrderParams,
1356 ) -> Result<DeriveOrder> {
1357 let reservation = self
1358 .reserve_matching_request(
1359 methods::PRIVATE_TRIGGER_ORDER,
1360 ¶ms.order.instrument_name,
1361 )
1362 .await?;
1363 self.submit_trigger_order_after_rate_limit(params, reservation)
1364 .await
1365 .map(|result| result.order)
1366 }
1367
1368 pub(crate) async fn submit_trigger_order_after_rate_limit(
1369 &self,
1370 params: &DeriveTriggerOrderParams,
1371 reservation: MatchingRateLimitReservation,
1372 ) -> Result<DeriveOrderResult> {
1373 self.submit_order_after_rate_limit(¶ms.order, reservation)
1374 .await
1375 }
1376
1377 pub async fn modify_order(&self, params: &DeriveReplaceParams) -> Result<DeriveReplaceOutcome> {
1385 let reservation = self
1386 .reserve_matching_request(methods::PRIVATE_REPLACE, ¶ms.order.instrument_name)
1387 .await?;
1388 self.modify_order_after_rate_limit(params, reservation)
1389 .await
1390 .and_then(|response| {
1391 response
1392 .result
1393 .and_then(|result| {
1394 result.into_outcome(
1395 ¶ms.order_id_to_cancel,
1396 ¶ms.order.label,
1397 params.order.envelope.subaccount_id,
1398 params.order.instrument_name.as_str(),
1399 )
1400 })
1401 .map_err(|message| {
1402 DeriveWsError::Serde(<serde_json::Error as serde::de::Error>::custom(
1403 message,
1404 ))
1405 })
1406 })
1407 }
1408
1409 pub(crate) async fn modify_order_after_rate_limit(
1410 &self,
1411 params: &DeriveReplaceParams,
1412 reservation: MatchingRateLimitReservation,
1413 ) -> Result<DeriveReplaceResponse> {
1414 self.ensure_authenticated(methods::PRIVATE_REPLACE)?;
1415 debug_assert_eq!(reservation.method, methods::PRIVATE_REPLACE);
1416 let cmd_tx = self.cmd_tx.read().await.clone();
1417 self.ensure_authenticated(reservation.method)?;
1418 self.commit_matching_reservation(reservation)?;
1419 let response = send_raw_after_rate_limit(
1420 &self.rate_limiter,
1421 &cmd_tx,
1422 methods::PRIVATE_REPLACE,
1423 params,
1424 self.request_timeout,
1425 )
1426 .await?;
1427 Ok(DeriveReplaceResponse::from_raw(&response))
1428 }
1429
1430 pub async fn cancel_order(&self, params: &DeriveCancelParams) -> Result<()> {
1436 self.require_authenticated(methods::PRIVATE_CANCEL).await?;
1437 let cmd_tx = self.cmd_tx.read().await.clone();
1438 let _: DeriveEmptyResult = send_request_for_instrument(
1439 &self.rate_limiter,
1440 &cmd_tx,
1441 methods::PRIVATE_CANCEL,
1442 params,
1443 self.request_timeout,
1444 params.instrument_name,
1445 )
1446 .await?;
1447 Ok(())
1448 }
1449
1450 pub async fn cancel_by_instrument(
1456 &self,
1457 params: &DeriveCancelByInstrumentParams,
1458 ) -> Result<DeriveCancelByInstrumentResult> {
1459 self.require_authenticated(methods::PRIVATE_CANCEL_BY_INSTRUMENT)
1460 .await?;
1461 let cmd_tx = self.cmd_tx.read().await.clone();
1462 send_request_typed_for_instrument(
1463 &self.rate_limiter,
1464 &cmd_tx,
1465 methods::PRIVATE_CANCEL_BY_INSTRUMENT,
1466 params,
1467 self.request_timeout,
1468 params.instrument_name,
1469 )
1470 .await
1471 }
1472
1473 pub async fn cancel_trigger_order(
1479 &self,
1480 params: &DeriveCancelTriggerOrderParams,
1481 ) -> Result<DeriveOrder> {
1482 self.require_authenticated(methods::PRIVATE_CANCEL_TRIGGER_ORDER)
1483 .await?;
1484 let cmd_tx = self.cmd_tx.read().await.clone();
1485 send_request_typed(
1486 &self.rate_limiter,
1487 &cmd_tx,
1488 methods::PRIVATE_CANCEL_TRIGGER_ORDER,
1489 params,
1490 self.request_timeout,
1491 )
1492 .await
1493 }
1494
1495 pub(crate) async fn cancel_trigger_order_current(
1496 &self,
1497 trigger: &DeriveCancelTriggerOrderParams,
1498 regular: &DeriveCancelParams,
1499 is_active: impl Fn() -> bool,
1500 ) -> Result<Option<DeriveOrder>> {
1501 self.require_authenticated(methods::PRIVATE_CANCEL_TRIGGER_ORDER)
1502 .await?;
1503 self.rate_limiter
1504 .await_class_ready(RateClass::NonMatching, None)
1505 .await;
1506 let cmd_tx = self.cmd_tx.read().await.clone();
1507
1508 if is_active() {
1509 return self.cancel_order(regular).await.map(|()| None);
1510 }
1511
1512 self.ensure_authenticated(methods::PRIVATE_CANCEL_TRIGGER_ORDER)?;
1513 send_request_typed_after_rate_limit(
1514 &self.rate_limiter,
1515 &cmd_tx,
1516 methods::PRIVATE_CANCEL_TRIGGER_ORDER,
1517 trigger,
1518 self.request_timeout,
1519 )
1520 .await
1521 .map(Some)
1522 }
1523
1524 pub async fn cancel_by_label(
1531 &self,
1532 params: &DeriveCancelByLabelParams,
1533 ) -> Result<DeriveCancelByLabelResult> {
1534 self.require_authenticated(methods::PRIVATE_CANCEL_BY_LABEL)
1535 .await?;
1536 let cmd_tx = self.cmd_tx.read().await.clone();
1537 send_request_typed(
1538 &self.rate_limiter,
1539 &cmd_tx,
1540 methods::PRIVATE_CANCEL_BY_LABEL,
1541 params,
1542 self.request_timeout,
1543 )
1544 .await
1545 }
1546
1547 pub async fn get_trigger_orders(
1554 &self,
1555 params: &DeriveGetTriggerOrdersParams,
1556 ) -> Result<DeriveOpenOrdersResult> {
1557 self.require_authenticated(methods::PRIVATE_GET_TRIGGER_ORDERS)
1558 .await?;
1559 let cmd_tx = self.cmd_tx.read().await.clone();
1560 send_request_typed(
1561 &self.rate_limiter,
1562 &cmd_tx,
1563 methods::PRIVATE_GET_TRIGGER_ORDERS,
1564 params,
1565 self.request_timeout,
1566 )
1567 .await
1568 }
1569
1570 pub async fn cancel_all_orders(&self, params: &DeriveCancelAllParams) -> Result<()> {
1577 self.require_authenticated(methods::PRIVATE_CANCEL_ALL)
1578 .await?;
1579 let cmd_tx = self.cmd_tx.read().await.clone();
1580 let _: DeriveEmptyResult = send_request(
1581 &self.rate_limiter,
1582 &cmd_tx,
1583 methods::PRIVATE_CANCEL_ALL,
1584 params,
1585 self.request_timeout,
1586 )
1587 .await?;
1588 Ok(())
1589 }
1590
1591 pub(crate) async fn reserve_matching_request(
1592 &self,
1593 operation: &'static str,
1594 instrument_name: &Ustr,
1595 ) -> Result<MatchingRateLimitReservation> {
1596 self.require_authenticated(operation).await?;
1597 debug_assert_eq!(rate_class_for_method(operation), RateClass::Matching);
1598 let quota = self
1599 .rate_limiter
1600 .reserve_matching_ready(instrument_name)
1601 .await;
1602 self.ensure_authenticated(operation)?;
1603 Ok(MatchingRateLimitReservation {
1604 method: operation,
1605 quota,
1606 })
1607 }
1608
1609 fn commit_matching_reservation(&self, reservation: MatchingRateLimitReservation) -> Result<()> {
1610 reservation
1611 .quota
1612 .commit()
1613 .map_err(|retry_after| DeriveWsError::RateLimited { retry_after })
1614 }
1615
1616 fn ensure_authenticated(&self, operation: &'static str) -> Result<()> {
1617 if self.auth_tracker.is_authenticated() {
1618 return Ok(());
1619 }
1620
1621 Err(DeriveWsError::Authentication {
1622 operation: operation.to_string(),
1623 reason: "WebSocket session is not authenticated".to_string(),
1624 })
1625 }
1626
1627 async fn require_authenticated(&self, operation: &'static str) -> Result<()> {
1628 if self
1629 .auth_tracker
1630 .wait_for_authenticated(self.request_timeout)
1631 .await
1632 {
1633 return Ok(());
1634 }
1635
1636 Err(DeriveWsError::Authentication {
1637 operation: operation.to_string(),
1638 reason: "WebSocket session is not authenticated".to_string(),
1639 })
1640 }
1641}
1642
1643#[derive(Debug, Clone, Copy, PartialEq, Eq)]
1647enum RequestRateLimit {
1648 Await(Option<Ustr>),
1651 Reserved,
1654}
1655
1656async fn send_raw<P>(
1657 rate_limiter: &WsRateLimiter,
1658 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1659 method: &'static str,
1660 params: &P,
1661 timeout: Duration,
1662) -> Result<Box<RawValue>>
1663where
1664 P: Serialize + ?Sized,
1665{
1666 send_raw_with_rate_limit(
1667 rate_limiter,
1668 cmd_tx,
1669 method,
1670 params,
1671 timeout,
1672 RequestRateLimit::Await(None),
1673 None,
1674 )
1675 .await
1676}
1677
1678async fn send_raw_for_instrument<P>(
1681 rate_limiter: &WsRateLimiter,
1682 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1683 method: &'static str,
1684 params: &P,
1685 timeout: Duration,
1686 instrument_name: Ustr,
1687) -> Result<Box<RawValue>>
1688where
1689 P: Serialize + ?Sized,
1690{
1691 send_raw_with_rate_limit(
1692 rate_limiter,
1693 cmd_tx,
1694 method,
1695 params,
1696 timeout,
1697 RequestRateLimit::Await(Some(instrument_name)),
1698 None,
1699 )
1700 .await
1701}
1702
1703async fn send_raw_after_rate_limit<P>(
1704 rate_limiter: &WsRateLimiter,
1705 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1706 method: &'static str,
1707 params: &P,
1708 timeout: Duration,
1709) -> Result<Box<RawValue>>
1710where
1711 P: Serialize + ?Sized,
1712{
1713 send_raw_with_rate_limit(
1714 rate_limiter,
1715 cmd_tx,
1716 method,
1717 params,
1718 timeout,
1719 RequestRateLimit::Reserved,
1720 None,
1721 )
1722 .await
1723}
1724
1725async fn send_raw_with_rate_limit<P>(
1726 rate_limiter: &WsRateLimiter,
1727 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1728 method: &'static str,
1729 params: &P,
1730 timeout: Duration,
1731 rate_limit: RequestRateLimit,
1732 connection_epoch: Option<u64>,
1733) -> Result<Box<RawValue>>
1734where
1735 P: Serialize + ?Sized,
1736{
1737 let params = SecretString::from(serde_json::to_string(params)?);
1738
1739 if let RequestRateLimit::Await(instrument_name) = rate_limit {
1740 rate_limiter
1741 .await_class_ready(rate_class_for_method(method), instrument_name.as_ref())
1742 .await;
1743 }
1744
1745 let (response_tx, response_rx) = tokio::sync::oneshot::channel();
1746 cmd_tx
1747 .send(HandlerCommand::Request {
1748 method,
1749 params,
1750 connection_epoch,
1751 cancellation: None,
1752 response_tx,
1753 })
1754 .map_err(|_| DeriveWsError::NotConnected)?;
1755
1756 match time::timeout(timeout, response_rx).await {
1760 Ok(Ok(outcome)) => outcome,
1761 Ok(Err(_)) => Err(DeriveWsError::RequestCancelled {
1762 method: method.to_owned(),
1763 }),
1764 Err(_) => Err(DeriveWsError::Timeout {
1765 method: method.to_owned(),
1766 }),
1767 }
1768}
1769
1770async fn send_request<P, R>(
1773 rate_limiter: &WsRateLimiter,
1774 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1775 method: &'static str,
1776 params: &P,
1777 timeout: Duration,
1778) -> Result<R>
1779where
1780 P: Serialize + ?Sized,
1781 R: Default + DeserializeOwned,
1782{
1783 let value = send_raw(rate_limiter, cmd_tx, method, params, timeout).await?;
1784 decode_default_result(&value)
1785}
1786
1787async fn send_request_for_instrument<P, R>(
1790 rate_limiter: &WsRateLimiter,
1791 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1792 method: &'static str,
1793 params: &P,
1794 timeout: Duration,
1795 instrument_name: Ustr,
1796) -> Result<R>
1797where
1798 P: Serialize + ?Sized,
1799 R: Default + DeserializeOwned,
1800{
1801 let value = send_raw_for_instrument(
1802 rate_limiter,
1803 cmd_tx,
1804 method,
1805 params,
1806 timeout,
1807 instrument_name,
1808 )
1809 .await?;
1810 decode_default_result(&value)
1811}
1812
1813async fn send_request_typed_for_instrument<P, R>(
1815 rate_limiter: &WsRateLimiter,
1816 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1817 method: &'static str,
1818 params: &P,
1819 timeout: Duration,
1820 instrument_name: Ustr,
1821) -> Result<R>
1822where
1823 P: Serialize + ?Sized,
1824 R: DeserializeOwned,
1825{
1826 let value = send_raw_for_instrument(
1827 rate_limiter,
1828 cmd_tx,
1829 method,
1830 params,
1831 timeout,
1832 instrument_name,
1833 )
1834 .await?;
1835 Ok(serde_json::from_str(value.get())?)
1836}
1837
1838fn decode_default_result<R>(value: &RawValue) -> Result<R>
1839where
1840 R: Default + DeserializeOwned,
1841{
1842 if value.get() == "null" {
1843 Ok(R::default())
1844 } else {
1845 Ok(serde_json::from_str(value.get())?)
1846 }
1847}
1848
1849async fn send_request_on_connection<P, R>(
1850 rate_limiter: &WsRateLimiter,
1851 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1852 method: &'static str,
1853 params: &P,
1854 timeout: Duration,
1855 connection_epoch: u64,
1856) -> Result<R>
1857where
1858 P: Serialize + ?Sized,
1859 R: Default + DeserializeOwned,
1860{
1861 let value = send_raw_with_rate_limit(
1862 rate_limiter,
1863 cmd_tx,
1864 method,
1865 params,
1866 timeout,
1867 RequestRateLimit::Await(None),
1868 Some(connection_epoch),
1869 )
1870 .await?;
1871
1872 decode_default_result(&value)
1873}
1874
1875async fn send_request_typed<P, R>(
1878 rate_limiter: &WsRateLimiter,
1879 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1880 method: &'static str,
1881 params: &P,
1882 timeout: Duration,
1883) -> Result<R>
1884where
1885 P: Serialize + ?Sized,
1886 R: DeserializeOwned,
1887{
1888 let value = send_raw(rate_limiter, cmd_tx, method, params, timeout).await?;
1889 Ok(serde_json::from_str(value.get())?)
1890}
1891
1892async fn send_request_typed_after_rate_limit<P, R>(
1893 rate_limiter: &WsRateLimiter,
1894 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1895 method: &'static str,
1896 params: &P,
1897 timeout: Duration,
1898) -> Result<R>
1899where
1900 P: Serialize + ?Sized,
1901 R: DeserializeOwned,
1902{
1903 let value = send_raw_after_rate_limit(rate_limiter, cmd_tx, method, params, timeout).await?;
1904 Ok(serde_json::from_str(value.get())?)
1905}
1906
1907fn channel_topics(channels: &[DeriveWsChannel]) -> Vec<String> {
1908 channels.iter().map(ToString::to_string).collect()
1909}
1910
1911#[expect(
1912 clippy::too_many_arguments,
1913 reason = "authentication state is passed explicitly for epoch fencing"
1914)]
1915async fn login_via_handler(
1916 rate_limiter: &WsRateLimiter,
1917 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1918 auth_tracker: &AuthTracker,
1919 authenticated_epoch: &AtomicU64,
1920 connection_mode: &AtomicU8,
1921 connection_epoch: &AtomicU64,
1922 creds: &DeriveWsCredentials,
1923 timeout: Duration,
1924) -> Result<()> {
1925 let _receiver = auth_tracker.begin();
1926 let expected_epoch = connection_epoch.load(Ordering::Acquire);
1927
1928 match send_login_request(rate_limiter, cmd_tx, creds, timeout, expected_epoch).await {
1929 Ok(())
1930 if complete_session_authentication(
1931 auth_tracker,
1932 authenticated_epoch,
1933 connection_mode,
1934 connection_epoch,
1935 expected_epoch,
1936 ) =>
1937 {
1938 log::debug!("Derive WebSocket authenticated");
1939
1940 Ok(())
1941 }
1942 Ok(()) => {
1943 let e = DeriveWsError::transport(
1944 "connection changed while completing WebSocket authentication",
1945 );
1946 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1947 auth_tracker.fail(e.to_string());
1948 Err(e)
1949 }
1950 Err(e) => {
1951 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
1952 auth_tracker.fail(e.to_string());
1953 Err(e)
1954 }
1955 }
1956}
1957
1958async fn send_login_request(
1959 rate_limiter: &WsRateLimiter,
1960 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
1961 creds: &DeriveWsCredentials,
1962 timeout: Duration,
1963 connection_epoch: u64,
1964) -> Result<()> {
1965 rate_limiter.await_public_ready().await;
1966 let login = build_ws_login(&creds.wallet_address, &creds.signer)?;
1967
1968 let params = Zeroizing::new(WsLoginParams {
1969 wallet: login.wallet,
1970 timestamp: login.timestamp,
1971 signature: login.signature,
1972 });
1973
1974 let value = send_raw_with_rate_limit(
1975 rate_limiter,
1976 cmd_tx,
1977 methods::PUBLIC_LOGIN,
1978 &*params,
1979 timeout,
1980 RequestRateLimit::Reserved,
1981 Some(connection_epoch),
1982 )
1983 .await?;
1984
1985 let _: WsLoginResult =
1986 serde_json::from_str(value.get()).map_err(|e| DeriveWsError::Authentication {
1987 operation: methods::PUBLIC_LOGIN.to_string(),
1988 reason: format!("invalid login result: {e}"),
1989 })?;
1990
1991 Ok(())
1992}
1993
1994#[expect(
1995 clippy::too_many_arguments,
1996 reason = "recovery state is passed explicitly for epoch fencing"
1997)]
1998async fn recover_session(
1999 rate_limiter: &WsRateLimiter,
2000 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
2001 auth_tracker: &AuthTracker,
2002 authenticated_epoch: &AtomicU64,
2003 connection_mode: &AtomicU8,
2004 connection_epoch: &AtomicU64,
2005 creds: Option<&DeriveWsCredentials>,
2006 subscriptions: &SubscriptionState,
2007 subscription_lock: &tokio::sync::Mutex<()>,
2008 timeout: Duration,
2009) -> Result<u64> {
2010 let _guard = subscription_lock.lock().await;
2011 let mut channels = subscriptions.reset_after_reconnect();
2012 channels.sort();
2013 let _receiver = creds.map(|_| auth_tracker.begin());
2014
2015 for attempt in 1..=MAX_SESSION_RECOVERY_ATTEMPTS {
2016 let expected_epoch = wait_for_session_connection(connection_mode, connection_epoch).await?;
2017
2018 let result = recover_session_attempt(
2019 rate_limiter,
2020 cmd_tx,
2021 connection_mode,
2022 connection_epoch,
2023 creds,
2024 subscriptions,
2025 channels.clone(),
2026 timeout,
2027 expected_epoch,
2028 )
2029 .await;
2030
2031 match result {
2032 Ok(()) => {
2033 if creds.is_some()
2034 && !complete_session_authentication(
2035 auth_tracker,
2036 authenticated_epoch,
2037 connection_mode,
2038 connection_epoch,
2039 expected_epoch,
2040 )
2041 {
2042 continue;
2043 }
2044
2045 if creds.is_some() {
2046 log::info!("Derive WebSocket session re-authenticated");
2047 }
2048
2049 return Ok(expected_epoch);
2050 }
2051 Err(e) if attempt < MAX_SESSION_RECOVERY_ATTEMPTS && should_retry_ws_error(&e) => {
2052 let multiplier = 1_u32 << (attempt - 1);
2053 let delay = RECONNECT_BASE_BACKOFF
2054 .saturating_mul(multiplier)
2055 .min(RECONNECT_MAX_BACKOFF);
2056 log::warn!(
2057 "Derive WebSocket session recovery attempt {attempt}/{MAX_SESSION_RECOVERY_ATTEMPTS} failed: {}; retrying in {delay:?}",
2058 e.redacted_message(),
2059 );
2060 time::sleep(delay).await;
2061 }
2062 Err(e) => {
2063 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
2064
2065 if creds.is_some() {
2066 auth_tracker.fail(e.to_string());
2067 }
2068
2069 return Err(e);
2070 }
2071 }
2072 }
2073
2074 let e = DeriveWsError::transport("WebSocket session changed while recovery completed");
2075 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
2076
2077 if creds.is_some() {
2078 auth_tracker.fail(e.to_string());
2079 }
2080
2081 Err(e)
2082}
2083
2084#[expect(
2085 clippy::too_many_arguments,
2086 reason = "a session attempt retains its fixed connection epoch and native replay context"
2087)]
2088async fn recover_session_attempt(
2089 rate_limiter: &WsRateLimiter,
2090 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
2091 connection_mode: &AtomicU8,
2092 connection_epoch: &AtomicU64,
2093 creds: Option<&DeriveWsCredentials>,
2094 subscriptions: &SubscriptionState,
2095 channels: Vec<String>,
2096 timeout: Duration,
2097 expected_epoch: u64,
2098) -> Result<()> {
2099 if !session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
2100 return Err(DeriveWsError::transport(
2101 "connection changed before WebSocket session recovery",
2102 ));
2103 }
2104
2105 if let Some(creds) = creds {
2106 send_login_request(rate_limiter, cmd_tx, creds, timeout, expected_epoch).await?;
2107 }
2108
2109 subscribe_via_handler(
2110 rate_limiter,
2111 cmd_tx,
2112 channels,
2113 timeout,
2114 expected_epoch,
2115 subscriptions,
2116 )
2117 .await?;
2118
2119 if !session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
2120 return Err(DeriveWsError::transport(
2121 "connection changed during WebSocket session recovery",
2122 ));
2123 }
2124
2125 Ok(())
2126}
2127
2128async fn wait_for_session_connection(
2129 connection_mode: &AtomicU8,
2130 connection_epoch: &AtomicU64,
2131) -> Result<u64> {
2132 loop {
2133 match ConnectionMode::from_atomic(connection_mode) {
2134 ConnectionMode::Active => return Ok(connection_epoch.load(Ordering::Acquire)),
2135 ConnectionMode::Reconnect => time::sleep(RECONNECT_BASE_BACKOFF).await,
2136 ConnectionMode::Disconnect | ConnectionMode::Closed => {
2137 return Err(DeriveWsError::transport(
2138 "WebSocket closed during session recovery",
2139 ));
2140 }
2141 }
2142 }
2143}
2144
2145fn complete_session_authentication(
2146 auth_tracker: &AuthTracker,
2147 authenticated_epoch: &AtomicU64,
2148 connection_mode: &AtomicU8,
2149 connection_epoch: &AtomicU64,
2150 expected_epoch: u64,
2151) -> bool {
2152 if !session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
2153 return false;
2154 }
2155
2156 authenticated_epoch.store(expected_epoch, Ordering::Release);
2157 auth_tracker.succeed();
2158
2159 if session_connection_is_active(connection_mode, connection_epoch, expected_epoch) {
2160 true
2161 } else {
2162 authenticated_epoch.store(UNAUTHENTICATED_CONNECTION_EPOCH, Ordering::Release);
2163 auth_tracker.invalidate();
2164 false
2165 }
2166}
2167
2168fn session_connection_is_active(
2169 connection_mode: &AtomicU8,
2170 connection_epoch: &AtomicU64,
2171 expected_epoch: u64,
2172) -> bool {
2173 ConnectionMode::from_atomic(connection_mode).is_active()
2174 && connection_epoch.load(Ordering::Acquire) == expected_epoch
2175}
2176
2177async fn subscribe_via_handler(
2178 rate_limiter: &WsRateLimiter,
2179 cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
2180 channels: Vec<String>,
2181 timeout: Duration,
2182 connection_epoch: u64,
2183 subscriptions: &SubscriptionState,
2184) -> Result<()> {
2185 if channels.is_empty() {
2186 return Ok(());
2187 }
2188
2189 let params = WsSubscribeParams {
2190 channels: channels.into_iter().map(DeriveWsChannel::from).collect(),
2191 };
2192
2193 let result: WsSubscribeResult = send_request_on_connection(
2194 rate_limiter,
2195 cmd_tx,
2196 methods::PUBLIC_SUBSCRIBE,
2197 ¶ms,
2198 timeout,
2199 connection_epoch,
2200 )
2201 .await?;
2202
2203 subscription_outcome(
2204 ¶ms.channels,
2205 &result,
2206 SUBSCRIPTION_REPLAY_ACCEPTED_STATUSES,
2207 subscriptions,
2208 )
2209}
2210
2211fn subscription_outcome(
2212 requested: &[DeriveWsChannel],
2213 result: &WsSubscribeResult,
2214 accepted_statuses: &[&str],
2215 subscriptions: &SubscriptionState,
2216) -> Result<()> {
2217 let mut failures = Vec::new();
2218
2219 for channel in requested {
2220 let topic = channel.to_string();
2221 match result.status.get(channel) {
2222 Some(status)
2223 if accepted_statuses.contains(&status.as_str())
2224 && result.channels.contains(channel) =>
2225 {
2226 subscriptions.confirm_subscribe(&topic);
2227 }
2228 Some(status) if accepted_statuses.contains(&status.as_str()) => {
2229 subscriptions.mark_failure(&topic);
2230 failures.push(format!("{topic}: not subscribed"));
2231 }
2232 other => {
2233 subscriptions.mark_failure(&topic);
2234 let reason = other.map_or("missing channel status", |status| status.as_str());
2235 failures.push(format!("{topic}: {reason}"));
2236 }
2237 }
2238 }
2239
2240 if failures.is_empty() {
2241 Ok(())
2242 } else {
2243 Err(DeriveWsError::Subscription {
2244 details: failures.join(", "),
2245 })
2246 }
2247}
2248
2249fn unsubscribe_outcome(
2250 requested: &[DeriveWsChannel],
2251 result: &WsUnsubscribeResult,
2252 subscriptions: &SubscriptionState,
2253) -> Result<()> {
2254 let mut failures = Vec::new();
2255
2256 for channel in requested {
2257 let topic = channel.to_string();
2258 match result.status.get(channel) {
2259 Some(status) if status == "ok" && !result.remaining_subscriptions.contains(channel) => {
2260 subscriptions.confirm_unsubscribe(&topic);
2261 }
2262 Some(status) if status == "ok" => failures.push(format!("{topic}: still subscribed")),
2263 Some(status) => failures.push(format!("{topic}: {status}")),
2264 None => failures.push(format!("{topic}: missing channel status")),
2265 }
2266 }
2267
2268 if failures.is_empty() {
2269 Ok(())
2270 } else {
2271 Err(DeriveWsError::Subscription {
2272 details: failures.join(", "),
2273 })
2274 }
2275}
2276
2277pub(crate) struct DeriveReplaceResponse {
2278 pub result: std::result::Result<DeriveReplaceResult, String>,
2279 pub trades: Vec<DeriveTrade>,
2280}
2281
2282impl DeriveReplaceResponse {
2283 fn from_raw(response: &RawValue) -> Self {
2284 #[derive(Deserialize)]
2285 struct Trades {
2286 #[serde(default)]
2287 trades: Vec<Box<RawValue>>,
2288 }
2289 let trades = serde_json::from_str::<Trades>(response.get())
2290 .map(|response| salvage_rows(response.trades))
2291 .unwrap_or_default();
2292 let result = serde_json::from_str(response.get()).map_err(|e| e.to_string());
2293 Self { result, trades }
2294 }
2295}
2296
2297#[cfg(test)]
2298mod tests {
2299 use std::{env, time::SystemTime};
2300
2301 use axum::{Json, Router, routing::post};
2302 use nautilus_common::live::dst::time::Instant;
2303 use nautilus_model::identifiers::InstrumentId;
2304 use rstest::rstest;
2305 use serde_json::Value;
2306
2307 use super::*;
2308 use crate::{
2309 common::{
2310 consts::REST_URL_TESTNET,
2311 rate_limit::{DERIVE_RATE_WINDOW_SECS, RateLimits, RequestLimiter},
2312 retry::{is_fatal_ws_error, is_write_outcome_ambiguous_ws},
2313 },
2314 http::{DeriveHttpClient, client::DeriveCredentials},
2315 };
2316
2317 #[rstest]
2318 #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
2319 #[case(
2320 "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
2321 "wss://stream.example.com:9443"
2322 )]
2323 fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
2324 let client = DeriveWebSocketClient::new(
2325 Some(url.to_string()),
2326 DeriveEnvironment::Testnet,
2327 TransportBackend::default(),
2328 None,
2329 );
2330 let debug = format!("{client:?}");
2331 assert!(debug.contains(&format!("url: {origin:?}")));
2332 assert!(!debug.contains("SYN-"));
2333 assert_eq!(client.url.expose(), url);
2334 }
2335
2336 #[rstest]
2337 #[case::subscription_lock(false)]
2338 #[case::rate_limit(true)]
2339 #[tokio::test]
2340 async fn test_queued_book_subscribe_uses_reconnected_generation(#[case] paced: bool) {
2341 let tracker = BookSyncTracker::default();
2342 let channel = "orderbook.ETH-USDC.1.10";
2343 tracker.record_subscription(channel, InstrumentId::from("ETH-USDC.DERIVE"));
2344 let client = DeriveWebSocketClient::new(
2345 None,
2346 DeriveEnvironment::Testnet,
2347 TransportBackend::default(),
2348 None,
2349 )
2350 .with_book_sync(tracker.clone(), Duration::from_secs(10));
2351 let mut handle = client.subscription_handle();
2352
2353 handle.rate_limiter = Arc::new(RequestLimiter::new(
2354 RateLimits::new(None, None),
2355 MonotonicClock,
2356 ));
2357 handle.connection_epoch.load().store(1, Ordering::Release);
2358 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2359 *handle.cmd_tx.write().await = cmd_tx;
2360 if paced {
2361 while handle
2362 .rate_limiter
2363 .try_class_ready(RateClass::NonMatching, None)
2364 .is_ok()
2365 {}
2366 }
2367
2368 let lock = Arc::clone(&handle.subscription_lock);
2369
2370 let guard = if paced { None } else { Some(lock.lock().await) };
2371 let request = handle.subscribe_book_gated(channel.to_owned());
2372 tokio::pin!(request);
2373 assert!(futures_util::poll!(&mut request).is_pending());
2374
2375 tracker.reset_on_reconnect(Duration::from_secs(10));
2376 handle.connection_epoch.load().store(2, Ordering::Release);
2377 drop(guard);
2378 let ack: Value =
2379 serde_json::from_str(include_str!("../../test_data/spot/ws_subscribe_ack.json"))
2380 .unwrap();
2381 tokio::time::timeout(Duration::from_secs(8), async {
2382 tokio::join!(async { request.await.unwrap() }, async {
2383 let HandlerCommand::Request {
2384 method,
2385 params,
2386 connection_epoch,
2387 cancellation,
2388 response_tx,
2389 } = cmd_rx.recv().await.unwrap()
2390 else {
2391 panic!("expected book subscribe request");
2392 };
2393 let params: Value = serde_json::from_str(params.expose_secret()).unwrap();
2394 response_tx
2395 .send(Ok(serde_json::value::to_raw_value(&ack["result"]).unwrap()))
2396 .unwrap();
2397
2398 assert_eq!(method, methods::PUBLIC_SUBSCRIBE);
2399 assert_eq!(params, serde_json::json!({"channels": [channel]}));
2400 assert_eq!(connection_epoch, Some(2));
2401 assert!(!cancellation.unwrap().is_cancelled());
2402 },);
2403 })
2404 .await
2405 .unwrap();
2406 }
2407
2408 #[rstest]
2409 #[tokio::test]
2410 async fn test_replay_ack_cannot_open_queued_book_replacement() {
2411 let tracker = BookSyncTracker::default();
2412 let channel = "orderbook.ETH-USDC.1.10";
2413 let now = Instant::now();
2414 tracker.record_subscription(channel, InstrumentId::from("ETH-USDC.DERIVE"));
2415 tracker.write(channel).unwrap().gate.open();
2416 assert!(tracker.record_snapshot(channel, 100, None, now));
2417 let recovery = tracker
2418 .recoveries(Duration::from_secs(1), now + Duration::from_secs(2))
2419 .pop()
2420 .unwrap()
2421 .2;
2422 assert!(recovery.begin_replacement());
2423 tracker.reset_on_reconnect(Duration::from_secs(10));
2424 let client = DeriveWebSocketClient::new(
2425 None,
2426 DeriveEnvironment::Testnet,
2427 TransportBackend::default(),
2428 None,
2429 )
2430 .with_book_sync(tracker.clone(), Duration::from_secs(10));
2431 let mut handle = client.subscription_handle();
2432
2433 handle.rate_limiter = Arc::new(RequestLimiter::new(
2434 RateLimits::new(None, None),
2435 MonotonicClock,
2436 ));
2437 handle.connection_epoch.load().store(2, Ordering::Release);
2438 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2439 *handle.cmd_tx.write().await = cmd_tx;
2440 let lock = Arc::clone(&handle.subscription_lock);
2441 let guard = lock.lock().await;
2442 let request = handle.replace_book_gated(
2443 channel.to_owned(),
2444 CancellationToken::new(),
2445 &recovery,
2446 recovery.gate.clone(),
2447 );
2448 tokio::pin!(request);
2449 assert!(futures_util::poll!(&mut request).is_pending());
2450 let replay = tracker.write(channel).unwrap();
2451 replay.gate.open();
2452 drop(guard);
2453 let ack: Value =
2454 serde_json::from_str(include_str!("../../test_data/spot/ws_subscribe_ack.json"))
2455 .unwrap();
2456 tokio::time::timeout(Duration::from_secs(1), async {
2457 tokio::join!(async { request.await.unwrap() }, async {
2458 for method in [methods::PUBLIC_UNSUBSCRIBE, methods::PUBLIC_SUBSCRIBE] {
2459 let HandlerCommand::Request {
2460 method: actual_method,
2461 params,
2462 connection_epoch,
2463 cancellation,
2464 response_tx,
2465 } = cmd_rx.recv().await.unwrap()
2466 else {
2467 panic!("expected book replacement request");
2468 };
2469 let current = tracker.write(channel).unwrap();
2470 assert!(current.gate.lock().is_closed());
2471 assert!(!tracker.record_snapshot(channel, 101, Some(¤t.generation), now));
2472 assert!(replay.generation.is_cancelled());
2473 assert_eq!(actual_method, method);
2474 assert_eq!(
2475 serde_json::from_str::<Value>(params.expose_secret()).unwrap(),
2476 serde_json::json!({"channels": [channel]}),
2477 );
2478 assert_eq!(connection_epoch, Some(2));
2479 assert!(!cancellation.unwrap().is_cancelled());
2480 let result = if method == methods::PUBLIC_SUBSCRIBE {
2481 current.gate.open();
2482 ack["result"].clone()
2483 } else {
2484 serde_json::json!({"status": {channel: "ok"}, "remaining_subscriptions": []})
2485 };
2486 response_tx
2487 .send(Ok(serde_json::value::to_raw_value(&result).unwrap()))
2488 .unwrap();
2489 }
2490 });
2491 })
2492 .await
2493 .unwrap();
2494
2495 assert!(tracker.record_snapshot(channel, 101, None, now));
2496 assert!(recovery.is_accepted());
2497 }
2498
2499 #[tokio::test]
2500 async fn test_expired_matching_reservation_has_local_failure_evidence() {
2501 let client = DeriveWebSocketClient::new(
2502 None,
2503 DeriveEnvironment::Testnet,
2504 TransportBackend::default(),
2505 None,
2506 );
2507 let handle = client.execution_handle();
2508 let quota = handle
2509 .rate_limiter
2510 .reserve_matching_ready(&Ustr::from("ETH-PERP"))
2511 .await;
2512 time::sleep(Duration::from_secs(DERIVE_RATE_WINDOW_SECS)).await;
2513 for _ in 0..DERIVE_RATE_WINDOW_SECS {
2514 handle
2515 .rate_limiter
2516 .await_class_ready(RateClass::Matching, None)
2517 .await;
2518 }
2519
2520 let reservation = MatchingRateLimitReservation {
2521 method: methods::PRIVATE_ORDER,
2522 quota,
2523 };
2524
2525 let error = handle.commit_matching_reservation(reservation).unwrap_err();
2526
2527 assert!(matches!(error, DeriveWsError::RateLimited { .. }));
2528 assert!(should_retry_ws_error(&error));
2529 assert!(!is_fatal_ws_error(&error));
2530 assert!(!is_write_outcome_ambiguous_ws(&error));
2531 }
2532
2533 #[rstest]
2534 fn test_replace_response_preserves_order_with_undecodable_trade() {
2535 let order: Value = serde_json::from_str(include_str!(
2536 "../../test_data/perps/http_order_eth_partially_filled.json"
2537 ))
2538 .unwrap();
2539 let trade: Value = serde_json::from_str(include_str!(
2540 "../../test_data/perps/http_private_trade_eth.json"
2541 ))
2542 .unwrap();
2543 let response = DeriveReplaceResponse::from_raw(
2544 &serde_json::value::to_raw_value(&serde_json::json!({
2545 "order": order,
2546 "trades": [{"trade_id": "bad-row"}, trade],
2547 }))
2548 .unwrap(),
2549 );
2550 let outcome = response
2551 .result
2552 .unwrap()
2553 .into_outcome("old-order", "alpha-strategy", 42, "ETH-PERP")
2554 .unwrap();
2555
2556 let DeriveReplaceOutcome::Replaced(order) = outcome else {
2557 panic!("expected replaced outcome");
2558 };
2559
2560 let trades = response.trades;
2561 assert_eq!(order.order_id, "abc-123");
2562 assert_eq!(order.label, "alpha-strategy");
2563 assert_eq!(trades.len(), 1);
2564 assert_eq!(trades[0].trade_id, "trade-xyz");
2565 assert_eq!(trades[0].order_id, "order-abc");
2566 }
2567
2568 const SESSION_KEY_HEX: &str =
2569 "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd";
2570
2571 #[rstest]
2572 #[case(Value::Null)]
2573 #[case(serde_json::json!("bad-nonce"))]
2574 #[case(serde_json::json!("18446744073709551616"))]
2575 fn test_replace_response_preserves_financial_rows_with_invalid_order_nonce(
2576 #[case] nonce: Value,
2577 ) {
2578 let mut order: Value = serde_json::from_str(include_str!(
2579 "../../test_data/perps/http_order_eth_partially_filled.json"
2580 ))
2581 .unwrap();
2582 let trade: DeriveTrade = serde_json::from_str(include_str!(
2583 "../../test_data/perps/http_private_trade_eth.json"
2584 ))
2585 .unwrap();
2586 order["nonce"] = nonce;
2587 let response = DeriveReplaceResponse::from_raw(
2588 &serde_json::value::to_raw_value(
2589 &serde_json::json!({"order": order, "trades": [trade]}),
2590 )
2591 .unwrap(),
2592 );
2593 assert!(response.result.is_err());
2594 assert_eq!(
2595 serde_json::to_value(response.trades).unwrap(),
2596 serde_json::to_value(vec![trade]).unwrap()
2597 );
2598 }
2599
2600 #[tokio::test]
2601 async fn test_http_and_websocket_share_matching_consumption() {
2602 let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2603 let address = listener.local_addr().unwrap();
2604
2605 let server = tokio::spawn(async move {
2606 axum::serve(
2607 listener,
2608 Router::new().route(
2609 "/private/cancel",
2610 post(|| async { Json(serde_json::json!({"result": true})) }),
2611 ),
2612 )
2613 .await
2614 .unwrap();
2615 });
2616
2617 let wallet = "0x000000000000000000000000000000000000ab01";
2618 let http = DeriveHttpClient::with_credentials(
2619 format!("http://{address}"),
2620 DeriveCredentials::new(wallet, SESSION_KEY_HEX).unwrap(),
2621 None,
2622 None,
2623 None,
2624 )
2625 .unwrap();
2626 let ws = DeriveWebSocketClient::with_credentials(
2627 Some(format!("ws://{address}/ws")),
2628 DeriveEnvironment::Testnet,
2629 TransportBackend::default(),
2630 None,
2631 DeriveWsCredentials::new(wallet.to_ascii_uppercase(), SESSION_KEY_HEX).unwrap(),
2632 None,
2633 None,
2634 );
2635
2636 for _ in 0..5 {
2637 let result: bool = http
2638 .send_private_once("private/cancel", &serde_json::json!({}))
2639 .await
2640 .unwrap();
2641 assert!(result);
2642 }
2643
2644 server.abort();
2645 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2646 let params = serde_json::json!({});
2647 let request = send_raw(
2648 &ws.rate_limiter,
2649 &cmd_tx,
2650 methods::PRIVATE_CANCEL,
2651 ¶ms,
2652 Duration::from_secs(1),
2653 );
2654 tokio::pin!(request);
2655 assert!(futures_util::poll!(&mut request).is_pending());
2656 assert!(matches!(
2657 cmd_rx.try_recv(),
2658 Err(tokio::sync::mpsc::error::TryRecvError::Empty)
2659 ));
2660 }
2661
2662 #[tokio::test]
2663 async fn test_login_builds_auth_timestamp_after_public_quota_wait() {
2664 let rate_limiter = RequestLimiter::new(RateLimits::new(None, None), MonotonicClock);
2665 for _ in 0..25 {
2666 rate_limiter.await_public_ready().await;
2667 }
2668
2669 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2670
2671 let captured = tokio::spawn(async move {
2672 let Some(HandlerCommand::Request {
2673 method,
2674 params,
2675 response_tx,
2676 ..
2677 }) = cmd_rx.recv().await
2678 else {
2679 panic!("login request must enqueue");
2680 };
2681
2682 assert_eq!(method, methods::PUBLIC_LOGIN);
2683 let params: Value = serde_json::from_str(params.expose_secret()).unwrap();
2684 let timestamp: u64 = params["timestamp"].as_str().unwrap().parse().unwrap();
2685 let now = SystemTime::now()
2686 .duration_since(SystemTime::UNIX_EPOCH)
2687 .unwrap()
2688 .as_millis();
2689 response_tx
2690 .send(Ok(serde_json::value::to_raw_value(&serde_json::json!([
2691 77
2692 ]))
2693 .unwrap()))
2694 .unwrap();
2695 now.saturating_sub(u128::from(timestamp))
2696 });
2697
2698 let started = tokio::time::Instant::now();
2699 send_login_request(
2700 &rate_limiter,
2701 &cmd_tx,
2702 &DeriveWsCredentials::new(
2703 "0x000000000000000000000000000000000000ab02",
2704 SESSION_KEY_HEX,
2705 )
2706 .unwrap(),
2707 Duration::from_secs(1),
2708 0,
2709 )
2710 .await
2711 .unwrap();
2712 assert!(started.elapsed() >= Duration::from_secs(4));
2713 assert!(
2714 captured.await.unwrap() < 1_000,
2715 "login timestamp must remain fresh after the pacing wait"
2716 );
2717 }
2718
2719 #[tokio::test]
2720 #[ignore = "requires Derive public testnet network access"]
2721 async fn test_live_testnet_unsubscribe_acknowledgement() {
2722 let mut ws = DeriveWebSocketClient::new(
2723 None,
2724 DeriveEnvironment::Testnet,
2725 TransportBackend::default(),
2726 None,
2727 );
2728 ws.connect().await.unwrap();
2729 let cmd_tx = ws.cmd_tx.read().await.clone();
2730 let params = serde_json::json!({"channels": ["ticker_slim.ETH-PERP.1000"]});
2731 let subscribe = send_raw(
2732 &ws.rate_limiter,
2733 &cmd_tx,
2734 methods::PUBLIC_SUBSCRIBE,
2735 ¶ms,
2736 Duration::from_secs(15),
2737 )
2738 .await;
2739 let unsubscribe = send_raw(
2740 &ws.rate_limiter,
2741 &cmd_tx,
2742 methods::PUBLIC_UNSUBSCRIBE,
2743 ¶ms,
2744 Duration::from_secs(15),
2745 )
2746 .await;
2747 ws.disconnect().await.unwrap();
2748 let subscribe: Value = serde_json::from_str(subscribe.unwrap().get()).unwrap();
2749 let unsubscribe: Value = serde_json::from_str(unsubscribe.unwrap().get()).unwrap();
2750 println!("Public testnet subscribe acknowledgement: {subscribe}");
2751 println!("Public testnet unsubscribe acknowledgement: {unsubscribe}");
2752 assert_eq!(subscribe["status"]["ticker_slim.ETH-PERP.1000"], "ok");
2753 assert_eq!(unsubscribe["status"]["ticker_slim.ETH-PERP.1000"], "ok");
2754 assert_eq!(
2755 unsubscribe["remaining_subscriptions"],
2756 serde_json::json!([])
2757 );
2758 }
2759
2760 #[tokio::test]
2761 #[ignore = "requires Derive testnet credentials and network access"]
2762 async fn test_live_testnet_request_budget_scopes() {
2763 let wallet = env::var("DERIVE_TESTNET_WALLET_ADDRESS").expect("testnet wallet configured");
2764 let key = env::var("DERIVE_TESTNET_SESSION_PRIVATE_KEY").expect("testnet key configured");
2765 let subaccount: u64 = env::var("DERIVE_TESTNET_SUBACCOUNT_ID")
2766 .expect("testnet subaccount configured")
2767 .parse()
2768 .unwrap();
2769 let http = DeriveHttpClient::with_credentials(
2770 REST_URL_TESTNET,
2771 DeriveCredentials::new(&wallet, &key).unwrap(),
2772 Some(15),
2773 None,
2774 None,
2775 )
2776 .unwrap();
2777 let mut ws = DeriveWebSocketClient::with_credentials(
2778 None,
2779 DeriveEnvironment::Testnet,
2780 TransportBackend::default(),
2781 None,
2782 DeriveWsCredentials::new(&wallet, &key).unwrap(),
2783 None,
2784 None,
2785 );
2786 ws.connect()
2787 .await
2788 .expect("testnet WebSocket login succeeds");
2789 let cmd_tx = ws.cmd_tx.read().await.clone();
2790 let before: Value = http
2791 .send_private("public/getRateLimits", &serde_json::json!({}))
2792 .await
2793 .expect("authenticated HTTP budget query succeeds");
2794 let ws_limits = send_raw(
2795 &ws.rate_limiter,
2796 &cmd_tx,
2797 "public/getRateLimits",
2798 &serde_json::json!({}),
2799 Duration::from_secs(15),
2800 )
2801 .await
2802 .expect("authenticated WebSocket budget query succeeds");
2803 println!("Testnet authenticated HTTP budgets: {before}");
2804 println!(
2805 "Testnet authenticated WebSocket budgets: {}",
2806 ws_limits.get()
2807 );
2808 let label = format!("nautilus-pacing-{}", UUID4::new());
2809
2810 for scoped in [true, false] {
2811 let mut params = serde_json::json!({"subaccount_id": subaccount, "label": label});
2812 if scoped {
2813 params["instrument_name"] = serde_json::json!("ETH-PERP");
2814 }
2815
2816 let result: Value = http
2817 .send_private_once("private/cancel_by_label", ¶ms)
2818 .await
2819 .expect("testnet cancellation of unique empty label succeeds");
2820 let budgets: Value = http
2821 .send_private("public/getRateLimits", &serde_json::json!({}))
2822 .await
2823 .unwrap();
2824 println!("Testnet cancel-by-label scoped={scoped}: {result}; budgets={budgets}");
2825 let ws_budgets = send_raw(
2826 &ws.rate_limiter,
2827 &cmd_tx,
2828 "public/getRateLimits",
2829 &serde_json::json!({}),
2830 Duration::from_secs(15),
2831 )
2832 .await
2833 .unwrap();
2834 println!(
2835 "Testnet wallet budgets after HTTP scoped={scoped}: {}",
2836 ws_budgets.get()
2837 );
2838 }
2839
2840 let started = tokio::time::Instant::now();
2841 let name = Ustr::from("ETH-PERP");
2842 let params = serde_json::json!({"subaccount_id": subaccount, "label": label, "instrument_name": name});
2843
2844 for _ in 0..6 {
2845 let result = send_raw_for_instrument(
2846 &ws.rate_limiter,
2847 &cmd_tx,
2848 methods::PRIVATE_CANCEL_BY_LABEL,
2849 ¶ms,
2850 Duration::from_secs(15),
2851 name,
2852 )
2853 .await
2854 .expect("paced testnet WebSocket scoped cancellation succeeds");
2855 let result: Value = serde_json::from_str(result.get()).unwrap();
2856 assert_eq!(result["cancelled_orders"], 0);
2857 }
2858
2859 let elapsed = started.elapsed();
2860 assert!(
2861 elapsed >= Duration::from_secs(1),
2862 "six same-instrument cancels must exceed the local five-token burst"
2863 );
2864 println!(
2865 "Testnet six scoped WebSocket cancellations pass with local pacing in {elapsed:?}"
2866 );
2867 let after = send_raw(
2868 &ws.rate_limiter,
2869 &cmd_tx,
2870 "public/getRateLimits",
2871 &serde_json::json!({}),
2872 Duration::from_secs(15),
2873 )
2874 .await
2875 .unwrap();
2876 println!(
2877 "Testnet wallet budgets after paced WebSocket requests: {}",
2878 after.get()
2879 );
2880 ws.disconnect().await.unwrap();
2881 ws.connect()
2882 .await
2883 .expect("testnet reconnect authenticates with preserved pacing state");
2884 ws.disconnect().await.unwrap();
2885 }
2886
2887 #[rstest]
2888 fn test_public_client_defaults_to_environment_url() {
2889 let client = DeriveWebSocketClient::new(
2890 None,
2891 DeriveEnvironment::Mainnet,
2892 TransportBackend::default(),
2893 None,
2894 );
2895 assert_eq!(client.url(), "wss://api.derive.xyz/v3/ws");
2896 assert!(!client.is_authenticated());
2897 assert!(!client.is_active());
2898 assert_eq!(client.subscription_count(), 0);
2899 }
2900
2901 #[tokio::test]
2902 async fn test_execution_auth_barrier_waits_for_authentication() {
2903 let client = DeriveWebSocketClient::with_credentials(
2904 None,
2905 DeriveEnvironment::Mainnet,
2906 TransportBackend::default(),
2907 None,
2908 DeriveWsCredentials::new(
2909 "0x000000000000000000000000000000000000aaaa",
2910 "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2911 )
2912 .unwrap(),
2913 None,
2914 None,
2915 );
2916 let execution = client.execution_handle();
2917 let auth_tracker = execution.auth_tracker.clone();
2918 let _receiver = auth_tracker.begin();
2919 let tracker_for_task = auth_tracker.clone();
2920
2921 get_runtime().spawn(async move {
2922 tokio::time::sleep(Duration::from_millis(10)).await;
2923 tracker_for_task.succeed();
2924 });
2925
2926 execution
2927 .require_authenticated(methods::PRIVATE_ORDER)
2928 .await
2929 .expect("barrier should wait for successful authentication");
2930 }
2931
2932 #[tokio::test]
2933 async fn test_execution_auth_barrier_fails_on_terminal_auth_failure() {
2934 let client = DeriveWebSocketClient::with_credentials(
2935 None,
2936 DeriveEnvironment::Mainnet,
2937 TransportBackend::default(),
2938 None,
2939 DeriveWsCredentials::new(
2940 "0x000000000000000000000000000000000000aaaa",
2941 "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2942 )
2943 .unwrap(),
2944 None,
2945 None,
2946 );
2947 let execution = client.execution_handle();
2948 let _receiver = execution.auth_tracker.begin();
2949 execution.auth_tracker.fail("bad signature");
2950
2951 let error = execution
2952 .require_authenticated(methods::PRIVATE_ORDER)
2953 .await
2954 .expect_err("terminal auth failure must reject private operations");
2955
2956 assert!(matches!(error, DeriveWsError::Authentication { .. }));
2957 }
2958
2959 #[tokio::test]
2960 async fn test_session_recovery_waits_for_reconnecting_transport() {
2961 let connection_mode = Arc::new(AtomicU8::new(ConnectionMode::Reconnect as u8));
2962 let connection_epoch = Arc::new(AtomicU64::new(1));
2963 let mode_for_task = Arc::clone(&connection_mode);
2964 let epoch_for_task = Arc::clone(&connection_epoch);
2965
2966 get_runtime().spawn(async move {
2967 tokio::time::sleep(Duration::from_millis(10)).await;
2968 epoch_for_task.store(2, Ordering::Release);
2969 mode_for_task.store(ConnectionMode::Active as u8, Ordering::Release);
2970 });
2971
2972 let epoch = tokio::time::timeout(
2973 Duration::from_secs(1),
2974 wait_for_session_connection(&connection_mode, &connection_epoch),
2975 )
2976 .await
2977 .expect("session recovery should resume after reconnect")
2978 .expect("active replacement connection should be accepted");
2979
2980 assert_eq!(epoch, 2);
2981 }
2982
2983 #[rstest]
2984 fn test_testnet_client_routes_to_v3_url() {
2985 let client = DeriveWebSocketClient::new(
2986 None,
2987 DeriveEnvironment::Testnet,
2988 TransportBackend::default(),
2989 None,
2990 );
2991 assert_eq!(client.url(), "wss://testnet.api.derive.xyz/v3/ws");
2992 }
2993
2994 #[rstest]
2995 fn test_credentials_constructor_parses_session_key() {
2996 let creds = DeriveWsCredentials::new(
2997 "0x000000000000000000000000000000000000aaaa",
2998 "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
2999 )
3000 .unwrap();
3001 assert!(creds.wallet_address.starts_with("0x"));
3002 let client = DeriveWebSocketClient::with_credentials(
3003 None,
3004 DeriveEnvironment::Testnet,
3005 TransportBackend::default(),
3006 None,
3007 creds,
3008 None,
3009 None,
3010 );
3011 assert_eq!(client.url(), "wss://testnet.api.derive.xyz/v3/ws");
3012 assert!(!client.is_authenticated());
3013 }
3014
3015 #[rstest]
3016 fn test_credentials_debug_redacts_signer() {
3017 let creds = DeriveWsCredentials::new(
3018 "0xWALLET",
3019 "0x2ae8be44db8a590d20bffbe3b6872df9b569147d3bf6801a35a28281a4816bbd",
3020 )
3021 .unwrap();
3022 let debug = format!("{creds:?}");
3023 assert!(debug.contains("redacted"));
3024 assert!(debug.contains("0xWALLET"));
3025 assert!(!debug.contains("2ae8be44"));
3026 }
3027
3028 #[rstest]
3029 fn test_credentials_constructor_rejects_invalid_session_key() {
3030 let err = DeriveWsCredentials::new("0xWALLET", "not-a-hex-key").unwrap_err();
3031 assert!(err.to_string().contains("invalid session key"));
3032 }
3033
3034 #[rstest]
3035 #[tokio::test]
3036 async fn test_send_raw_closed_command_queue_is_definitive() {
3037 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
3038 drop(cmd_rx);
3039 let rate_limiter: WsRateLimiter =
3040 RequestLimiter::new(RateLimits::new(None, None), MonotonicClock);
3041 let error = send_raw(
3042 &rate_limiter,
3043 &cmd_tx,
3044 methods::PRIVATE_REPLACE,
3045 &serde_json::json!({}),
3046 Duration::from_millis(50),
3047 )
3048 .await
3049 .unwrap_err();
3050 assert!(matches!(error, DeriveWsError::NotConnected));
3051 assert!(!crate::common::retry::is_write_outcome_ambiguous_ws(&error));
3052 }
3053
3054 #[rstest]
3055 #[tokio::test]
3056 async fn test_send_raw_times_out_when_no_response_arrives() {
3057 let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
3060 let rate_limiter: WsRateLimiter =
3061 RequestLimiter::new(RateLimits::new(None, None), MonotonicClock);
3062 let err = send_raw(
3063 &rate_limiter,
3064 &cmd_tx,
3065 methods::PRIVATE_ORDER,
3066 &serde_json::json!({}),
3067 Duration::from_millis(50),
3068 )
3069 .await
3070 .expect_err("must time out");
3071
3072 match err {
3073 DeriveWsError::Timeout { method } => assert_eq!(method, methods::PRIVATE_ORDER),
3074 other => panic!("expected Timeout, was {other:?}"),
3075 }
3076 }
3077
3078 #[rstest]
3079 #[tokio::test]
3080 async fn test_send_request_typed_rejects_null_result() {
3081 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
3085 tokio::spawn(async move {
3086 if let Some(HandlerCommand::Request { response_tx, .. }) = cmd_rx.recv().await {
3087 let _ = response_tx.send(Ok(RawValue::NULL.to_owned()));
3088 }
3089 });
3090
3091 let rate_limiter: WsRateLimiter =
3092 RequestLimiter::new(RateLimits::new(None, None), MonotonicClock);
3093 let result: Result<DeriveOrderResult> = send_request_typed(
3094 &rate_limiter,
3095 &cmd_tx,
3096 methods::PRIVATE_ORDER,
3097 &serde_json::json!({}),
3098 Duration::from_secs(1),
3099 )
3100 .await;
3101 assert!(matches!(result, Err(DeriveWsError::Serde(_))));
3102 }
3103
3104 #[rstest]
3105 #[case::pending(false)]
3106 #[case::fired_during_pacing(true)]
3107 #[tokio::test]
3108 async fn test_trigger_cancel_selects_endpoint_after_pacing(#[case] fires: bool) {
3109 let client = DeriveWebSocketClient::new(
3110 None,
3111 DeriveEnvironment::Testnet,
3112 TransportBackend::default(),
3113 None,
3114 );
3115 let mut handle = client.execution_handle();
3116 handle.auth_tracker.succeed();
3117
3118 handle.rate_limiter = Arc::new(RequestLimiter::new(
3119 RateLimits::new(None, None),
3120 MonotonicClock,
3121 ));
3122 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
3123 *handle.cmd_tx.write().await = cmd_tx;
3124 for _ in 0..crate::common::rate_limit::DERIVE_WEBSOCKET_NON_MATCHING_TPS
3125 * crate::common::rate_limit::DERIVE_RATE_BURST_MULTIPLIER
3126 {
3127 handle
3128 .rate_limiter
3129 .try_class_ready(RateClass::NonMatching, None)
3130 .unwrap();
3131 }
3132
3133 assert!(
3134 handle
3135 .rate_limiter
3136 .try_class_ready(RateClass::NonMatching, None)
3137 .is_err()
3138 );
3139 let active = AtomicBool::new(false);
3140 let trigger = DeriveCancelTriggerOrderParams::new(42, "paced-trigger");
3141 let regular = DeriveCancelParams::new(42, "ETH-PERP", "paced-trigger");
3142 let cancel = handle
3143 .cancel_trigger_order_current(&trigger, ®ular, || active.load(Ordering::Acquire));
3144 tokio::pin!(cancel);
3145 assert!(futures_util::poll!(cancel.as_mut()).is_pending());
3146 assert!(cmd_rx.try_recv().is_err());
3147 active.store(fires, Ordering::Release);
3148 let mut canceled: DeriveOrder = serde_json::from_str(include_str!(
3149 "../../test_data/perps/http_order_eth_partially_filled.json"
3150 ))
3151 .unwrap();
3152 canceled.order_id = "paced-trigger".into();
3153 canceled.order_status = crate::common::enums::DeriveOrderStatus::Cancelled;
3154 canceled.filled_amount = rust_decimal::Decimal::ZERO;
3155 canceled.average_price = rust_decimal::Decimal::ZERO;
3156 let expected = (!fires).then_some(canceled);
3157
3158 let response = async {
3159 let HandlerCommand::Request {
3160 method,
3161 params,
3162 response_tx,
3163 ..
3164 } = cmd_rx.recv().await.unwrap()
3165 else {
3166 panic!("expected one cancel request")
3167 };
3168
3169 let params: Value = serde_json::from_str(params.expose_secret()).unwrap();
3170
3171 let expected_method = if fires {
3172 methods::PRIVATE_CANCEL
3173 } else {
3174 methods::PRIVATE_CANCEL_TRIGGER_ORDER
3175 };
3176
3177 assert_eq!(method, expected_method);
3178 assert_eq!(params["subaccount_id"], serde_json::json!(42));
3179 assert_eq!(params["order_id"], serde_json::json!("paced-trigger"));
3180 assert_eq!(
3181 params.get("instrument_name"),
3182 fires
3183 .then_some(®ular.instrument_name)
3184 .map(|name| serde_json::json!(name))
3185 .as_ref()
3186 );
3187 response_tx
3188 .send(Ok(serde_json::value::to_raw_value(&expected).unwrap()))
3189 .unwrap();
3190 };
3191
3192 let (result, ()) = tokio::join!(cancel, response);
3193 assert_eq!(
3194 serde_json::to_value(result.unwrap()).unwrap(),
3195 serde_json::to_value(expected).unwrap()
3196 );
3197 assert!(cmd_rx.try_recv().is_err());
3198 }
3199
3200 #[rstest]
3201 #[tokio::test]
3202 async fn test_reserved_send_does_not_wait_for_or_consume_second_quota_cell() {
3203 let rate_limiter: WsRateLimiter =
3204 RequestLimiter::new(RateLimits::new(None, None), MonotonicClock);
3205
3206 for _ in 0..5 {
3207 rate_limiter
3208 .try_class_ready(RateClass::Matching, None)
3209 .expect("reservation consumes the matching window");
3210 }
3211
3212 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
3213 tokio::spawn(async move {
3214 if let Some(HandlerCommand::Request { response_tx, .. }) = cmd_rx.recv().await {
3215 let _ = response_tx.send(Ok(serde_json::value::to_raw_value(
3216 &serde_json::json!({"accepted": true}),
3217 )
3218 .unwrap()));
3219 }
3220 });
3221
3222 let response = tokio::time::timeout(
3223 Duration::from_millis(100),
3224 send_raw_after_rate_limit(
3225 &rate_limiter,
3226 &cmd_tx,
3227 methods::PRIVATE_ORDER,
3228 &serde_json::json!({}),
3229 Duration::from_secs(1),
3230 ),
3231 )
3232 .await
3233 .expect("reserved send must not wait for quota")
3234 .expect("reserved send succeeds");
3235
3236 assert_eq!(
3237 serde_json::from_str::<Value>(response.get()).unwrap(),
3238 serde_json::json!({"accepted": true}),
3239 );
3240 assert!(
3241 rate_limiter
3242 .try_class_ready(RateClass::Matching, None)
3243 .is_err(),
3244 "reserved send must not consume a second cell",
3245 );
3246 }
3247}