1use std::{
19 collections::HashMap,
20 sync::{
21 Arc,
22 atomic::{AtomicBool, AtomicU8, AtomicU64, Ordering},
23 },
24};
25
26use arc_swap::ArcSwap;
27#[cfg(test)]
28use nautilus_core::string::secret::REDACTED;
29use nautilus_core::{AtomicMap, string::secret::SecretString};
30use nautilus_live::{
31 SocketControl,
32 task::{TaskGroup, TaskShutdownError},
33};
34use nautilus_model::{
35 data::BarType,
36 enums::BarAggregation,
37 identifiers::{AccountId, ClientOrderId, InstrumentId, StrategyId, TraderId, VenueOrderId},
38 instruments::{Instrument, InstrumentAny},
39};
40use nautilus_network::{
41 http::create_standard_nautilus_headers,
42 mode::ConnectionMode,
43 websocket::{
44 AuthTracker, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
45 channel_message_handler,
46 },
47};
48use parking_lot::RwLock;
49use tokio_util::sync::CancellationToken;
50use ustr::Ustr;
51
52pub const KRAKEN_SPOT_WS_TOPIC_DELIMITER: char = ':';
56
57use super::{
58 enums::{KrakenWsChannel, KrakenWsMethod},
59 handler::{SpotFeedHandler, SpotHandlerCommand},
60 level_2::L2Depths,
61 messages::{KrakenSpotWsMessage, KrakenWsChannelParams, KrakenWsParams, KrakenWsRequest},
62};
63use crate::{
64 common::{
65 consts::{
66 KRAKEN_RATE_LIMIT_KEY_ORDER, KRAKEN_RATE_LIMIT_KEY_SUBSCRIPTION,
67 KRAKEN_SPOT_WS_ORDER_QUOTA, KRAKEN_SPOT_WS_SUBSCRIPTION_QUOTA,
68 },
69 parse::normalize_spot_symbol,
70 },
71 config::KrakenDataClientConfig,
72 http::{KrakenSpotHttpClient, spot::client::KRAKEN_SPOT_DEFAULT_RATE_LIMIT_PER_SECOND},
73 websocket::error::KrakenWsError,
74};
75
76const WS_PING_MSG: &str = r#"{"method":"ping"}"#;
77
78#[derive(Debug)]
80pub struct KrakenSpotWebSocketClient {
81 url: String,
82 config: KrakenDataClientConfig,
83 signal: Arc<AtomicBool>,
84 connection_mode: Arc<ArcSwap<AtomicU8>>,
85 cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand>>>,
86 out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<KrakenSpotWsMessage>>>,
87 handler_tasks: Arc<TaskGroup>,
88 connect_lock: Arc<tokio::sync::Mutex<()>>,
89 subscriptions: SubscriptionState,
90 subscription_payloads: Arc<tokio::sync::RwLock<HashMap<String, SecretString>>>,
91 auth_tracker: AuthTracker,
92 cancellation_token: CancellationToken,
93 req_id_counter: Arc<AtomicU64>,
94 auth_token: Arc<tokio::sync::RwLock<Option<SecretString>>>,
95 account_id: Arc<RwLock<Option<AccountId>>>,
96 truncated_id_map: Arc<AtomicMap<String, ClientOrderId>>,
97 instruments: Arc<AtomicMap<InstrumentId, InstrumentAny>>,
98 l2_depths: L2Depths,
99 l3_depths: Arc<parking_lot::Mutex<ahash::AHashMap<String, u32>>>,
100 transport_backend: TransportBackend,
101 proxy_url: Option<SecretString>,
102 socket_control: Option<SocketControl>,
103}
104
105impl Clone for KrakenSpotWebSocketClient {
106 fn clone(&self) -> Self {
107 Self {
108 url: self.url.clone(),
109 config: self.config.clone(),
110 signal: Arc::clone(&self.signal),
111 connection_mode: Arc::clone(&self.connection_mode),
112 cmd_tx: Arc::clone(&self.cmd_tx),
113 out_rx: self.out_rx.clone(),
114 handler_tasks: Arc::clone(&self.handler_tasks),
115 connect_lock: Arc::clone(&self.connect_lock),
116 subscriptions: self.subscriptions.clone(),
117 subscription_payloads: Arc::clone(&self.subscription_payloads),
118 auth_tracker: self.auth_tracker.clone(),
119 cancellation_token: self.cancellation_token.clone(),
120 req_id_counter: self.req_id_counter.clone(),
121 auth_token: self.auth_token.clone(),
122 account_id: Arc::clone(&self.account_id),
123 truncated_id_map: Arc::clone(&self.truncated_id_map),
124 instruments: Arc::clone(&self.instruments),
125 l2_depths: self.l2_depths.clone(),
126 l3_depths: Arc::clone(&self.l3_depths),
127 transport_backend: self.transport_backend,
128 proxy_url: self.proxy_url.clone(),
129 socket_control: self.socket_control.clone(),
130 }
131 }
132}
133
134impl KrakenSpotWebSocketClient {
135 pub fn new(
137 config: KrakenDataClientConfig,
138 cancellation_token: CancellationToken,
139 proxy_url: Option<String>,
140 ) -> Self {
141 let url = if config.ws_private_url.is_some() {
142 config.ws_private_url()
143 } else {
144 config.ws_public_url()
145 };
146 Self::new_with_url(url, config, cancellation_token, proxy_url)
147 }
148
149 pub fn l3(
154 config: KrakenDataClientConfig,
155 cancellation_token: CancellationToken,
156 proxy_url: Option<String>,
157 ) -> Self {
158 let url = config.ws_l3_url();
159 Self::new_with_url(url, config, cancellation_token, proxy_url)
160 }
161
162 fn new_with_url(
163 url: String,
164 mut config: KrakenDataClientConfig,
165 cancellation_token: CancellationToken,
166 proxy_url: Option<String>,
167 ) -> Self {
168 let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<SpotHandlerCommand>();
169 let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
170 let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
171
172 let transport_backend = config.transport_backend;
173 config.proxy_url = proxy_url.clone().map(SecretString::from);
174
175 Self {
176 url,
177 config,
178 signal: Arc::new(AtomicBool::new(false)),
179 connection_mode,
180 cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
181 out_rx: None,
182 handler_tasks: Arc::new(TaskGroup::new()),
183 connect_lock: Arc::new(tokio::sync::Mutex::new(())),
184 subscriptions: SubscriptionState::new(KRAKEN_SPOT_WS_TOPIC_DELIMITER),
185 subscription_payloads: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
186 auth_tracker: AuthTracker::new(),
187 cancellation_token,
188 req_id_counter: Arc::new(AtomicU64::new(0)),
189 auth_token: Arc::new(tokio::sync::RwLock::new(None)),
190 account_id: Arc::new(RwLock::new(None)),
191 truncated_id_map: Arc::new(AtomicMap::new()),
192 instruments: Arc::new(AtomicMap::new()),
193 l2_depths: L2Depths::default(),
194 l3_depths: Arc::new(parking_lot::Mutex::new(ahash::AHashMap::new())),
195 transport_backend,
196 proxy_url: proxy_url.map(SecretString::from),
197 socket_control: None,
198 }
199 }
200
201 pub(crate) fn begin_shutdown(&self) {
202 self.handler_tasks.begin_shutdown();
203 self.cancellation_token.cancel();
204 self.signal.store(true, Ordering::Relaxed);
205 }
206
207 #[must_use]
209 pub fn with_socket_control(mut self, control: SocketControl) -> Self {
210 self.socket_control = Some(control);
211 self
212 }
213
214 fn get_next_req_id(&self) -> u64 {
215 self.req_id_counter.fetch_add(1, Ordering::Relaxed) + 1
216 }
217
218 pub fn req_id_counter(&self) -> Arc<AtomicU64> {
220 self.req_id_counter.clone()
221 }
222
223 pub async fn handler_command_sender(
225 &self,
226 ) -> tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand> {
227 self.cmd_tx.read().await.clone()
228 }
229
230 pub fn handler_command_handle(
235 &self,
236 ) -> Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand>>> {
237 self.cmd_tx.clone()
238 }
239
240 pub async fn auth_token(&self) -> Option<SecretString> {
242 self.auth_token.read().await.clone()
243 }
244
245 pub fn auth_token_blocking(&self) -> Option<SecretString> {
251 self.auth_token
252 .try_read()
253 .ok()
254 .and_then(|guard| guard.clone())
255 }
256
257 pub(crate) fn auth_token_handle(&self) -> Arc<tokio::sync::RwLock<Option<SecretString>>> {
260 self.auth_token.clone()
261 }
262
263 pub async fn connect(&mut self) -> Result<(), KrakenWsError> {
265 let connect_lock = Arc::clone(&self.connect_lock);
266 let _connect_guard = connect_lock.lock().await;
267
268 log::debug!("Connecting to {}", self.url);
269
270 if !self.handler_tasks.is_open() || !self.handler_tasks.is_empty() {
271 self.disconnect_locked().await?;
272 self.handler_tasks.start_generation().map_err(|e| {
273 KrakenWsError::ConnectionError(format!(
274 "Failed to start WebSocket handler task generation: {e}"
275 ))
276 })?;
277 }
278 let handler_spawner = self.handler_tasks.spawner().map_err(|e| {
279 KrakenWsError::ConnectionError(format!(
280 "Failed to acquire WebSocket handler task spawner: {e}"
281 ))
282 })?;
283
284 if self.cancellation_token.is_cancelled() {
285 self.cancellation_token = CancellationToken::new();
286 }
287
288 self.signal.store(false, Ordering::Relaxed);
289
290 let (raw_handler, raw_rx) = channel_message_handler();
291 let headers = create_standard_nautilus_headers();
292
293 let ws_config = WebSocketConfig {
294 url: self.url.clone(),
295 headers,
296 heartbeat_interval_secs: Some(self.config.heartbeat_interval_secs),
297 heartbeat_payload: Some(WS_PING_MSG.to_string()),
298 connect_timeout_ms: Some(5_000),
299 reconnect_delay_initial_ms: Some(500),
300 reconnect_delay_max_ms: Some(5_000),
301 reconnect_backoff_factor: Some(1.5),
302 reconnect_jitter_ms: Some(250),
303 reconnect_max_attempts: None,
304 heartbeat_timeout_secs: None,
307 idle_timeout_ms: (self.config.ws_idle_timeout_ms != 0)
308 .then_some(self.config.ws_idle_timeout_ms),
309 backend: self.transport_backend,
310 proxy_url: self
311 .proxy_url
312 .as_ref()
313 .map(|value| value.expose_secret().to_owned()),
314 };
315
316 let keyed_quotas = vec![
317 (
318 KRAKEN_RATE_LIMIT_KEY_SUBSCRIPTION[0].to_string(),
319 *KRAKEN_SPOT_WS_SUBSCRIPTION_QUOTA,
320 ),
321 (
322 KRAKEN_RATE_LIMIT_KEY_ORDER[0].to_string(),
323 *KRAKEN_SPOT_WS_ORDER_QUOTA,
324 ),
325 ];
326
327 let ws_client = WebSocketClient::builder()
328 .config(ws_config)
329 .message_handler(raw_handler)
330 .keyed_quotas(keyed_quotas)
331 .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
332 .connect()
333 .await
334 .map_err(|e| KrakenWsError::ConnectionError(e.to_string()))?;
335
336 self.connection_mode
338 .store(ws_client.connection_mode_atomic());
339 let reconnect_handle = ws_client.reconnect_handle();
340
341 let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<KrakenSpotWsMessage>();
342 self.out_rx = Some(Arc::new(out_rx));
343
344 let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<SpotHandlerCommand>();
345 *self.cmd_tx.write().await = cmd_tx.clone();
346
347 if let Err(e) = cmd_tx.send(SpotHandlerCommand::SetClient(ws_client)) {
348 return Err(KrakenWsError::ConnectionError(format!(
349 "Failed to send WebSocketClient to handler: {e}"
350 )));
351 }
352
353 if let Some(control) = &self.socket_control {
354 control.register(move || reconnect_handle.request_reconnect());
355 }
356
357 let signal = self.signal.clone();
358 let subscriptions = self.subscriptions.clone();
359 let subscription_payloads = self.subscription_payloads.clone();
360 let config_for_reconnect = self.config.clone();
361 let auth_token_for_reconnect = self.auth_token.clone();
362 let auth_tracker_for_reconnect = self.auth_tracker.clone();
363 let cmd_tx_for_reconnect = cmd_tx.clone();
364
365 let handler_task = async move {
366 let mut handler =
367 SpotFeedHandler::new(signal.clone(), cmd_rx, raw_rx, subscriptions.clone());
368
369 loop {
370 match handler.next().await {
371 Some(KrakenSpotWsMessage::Reconnected) => {
372 if signal.load(Ordering::Relaxed) {
373 continue;
374 }
375 log::info!("WebSocket reconnected, resubscribing");
376
377 subscriptions.reset_after_reconnect();
378
379 let payloads = subscription_payloads.read().await;
380 if payloads.is_empty() {
381 log::debug!("No subscriptions to restore after reconnection");
382 } else {
383 let had_auth = auth_token_for_reconnect.read().await.is_some();
384
385 if had_auth && config_for_reconnect.has_api_credentials() {
386 log::debug!("Re-authenticating after reconnect");
387
388 auth_tracker_for_reconnect.invalidate();
389 let _rx = auth_tracker_for_reconnect.begin();
390
391 match refresh_auth_token(&config_for_reconnect).await {
392 Ok(new_token) => {
393 *auth_token_for_reconnect.write().await = Some(new_token);
394 auth_tracker_for_reconnect.succeed();
395 log::debug!("Re-authentication successful");
396 }
397 Err(e) => {
398 log::error!(
399 "Failed to re-authenticate after reconnect: {e}"
400 );
401 *auth_token_for_reconnect.write().await = None;
402 auth_tracker_for_reconnect.fail(e.to_string());
403 }
404 }
405 }
406
407 log::debug!(
408 "Resubscribing after reconnection: count={}",
409 payloads.len()
410 );
411
412 for (topic, payload) in payloads.iter() {
413 let needs_token =
414 topic == "executions" || topic.starts_with("level3:");
415 let payload = if needs_token {
416 let auth_token = auth_token_for_reconnect.read().await.clone();
417 match auth_token {
418 Some(token) => {
419 match update_auth_token_in_payload(
420 payload.expose_secret(),
421 token.expose_secret(),
422 ) {
423 Ok(p) => p,
424 Err(e) => {
425 log::error!("Failed to update auth token: {e}");
426 continue;
427 }
428 }
429 }
430 None => {
431 log::warn!(
432 "Cannot resubscribe to {topic}: no auth token"
433 );
434 continue;
435 }
436 }
437 } else {
438 payload.clone()
439 };
440
441 if let Err(e) = cmd_tx_for_reconnect
442 .send(SpotHandlerCommand::Subscribe { payload })
443 {
444 log::error!(
445 "Failed to send resubscribe command: error={e}, \
446 topic={topic}"
447 );
448 }
449
450 subscriptions.mark_subscribe(topic);
451 }
452 }
453
454 if out_tx.send(KrakenSpotWsMessage::Reconnected).is_err() {
455 if handler.is_stopped() {
456 log::debug!("Failed to send message (receiver dropped)");
457 } else {
458 log::error!("Failed to send message (receiver dropped)");
459 }
460 break;
461 }
462 }
463 Some(msg) => {
464 if out_tx.send(msg).is_err() {
465 if handler.is_stopped() {
466 log::debug!("Failed to send message (receiver dropped)");
467 } else {
468 log::error!("Failed to send message (receiver dropped)");
469 }
470 break;
471 }
472 }
473 None => {
474 if handler.is_stopped() {
475 log::debug!("Stop signal received, ending message processing");
476 break;
477 }
478 log::warn!("WebSocket stream ended unexpectedly");
479 break;
480 }
481 }
482 }
483
484 log::debug!("Handler task exiting");
485 };
486
487 if let Err(e) = handler_spawner.spawn(handler_task) {
488 if let Some(control) = &self.socket_control {
489 control.deregister();
490 }
491 self.out_rx = None;
492 return Err(KrakenWsError::ConnectionError(format!(
493 "Failed to register WebSocket handler task: {e}"
494 )));
495 }
496
497 log::debug!("WebSocket connected successfully");
498 Ok(())
499 }
500
501 pub async fn disconnect(&mut self) -> Result<(), KrakenWsError> {
503 let connect_lock = Arc::clone(&self.connect_lock);
504 let _connect_guard = connect_lock.lock().await;
505 self.disconnect_locked().await
506 }
507
508 async fn disconnect_locked(&self) -> Result<(), KrakenWsError> {
509 log::debug!("Disconnecting WebSocket");
510
511 self.handler_tasks.begin_shutdown();
512 self.signal.store(true, Ordering::Relaxed);
513
514 if let Err(e) = self
515 .cmd_tx
516 .read()
517 .await
518 .send(SpotHandlerCommand::Disconnect)
519 {
520 log::debug!(
521 "Failed to send disconnect command (handler may already be shut down): {e}"
522 );
523 }
524
525 let task_result = self
526 .handler_tasks
527 .finish_shutdown(
528 tokio::time::Duration::from_secs(2),
529 tokio::time::Duration::from_secs(2),
530 )
531 .await;
532
533 self.subscriptions.clear();
534 self.subscription_payloads.write().await.clear();
535 self.auth_tracker.fail("Disconnected");
536
537 self.l3_depths.lock().clear();
538 self.l2_depths.clear();
539
540 if let Some(control) = &self.socket_control {
541 control.deregister();
542 }
543
544 match task_result {
545 Ok(()) => Ok(()),
546 Err(error @ TaskShutdownError::Timeout { .. }) => Err(KrakenWsError::Timeout(format!(
547 "Spot WebSocket handler shutdown timed out: {error}"
548 ))),
549 Err(e) => Err(KrakenWsError::Disconnected(format!(
550 "Spot WebSocket handler shutdown failed: {e}"
551 ))),
552 }
553 }
554
555 pub async fn close(&mut self) -> Result<(), KrakenWsError> {
557 self.disconnect().await
558 }
559
560 pub async fn wait_until_active(&self, timeout_secs: f64) -> Result<(), KrakenWsError> {
562 let timeout = tokio::time::Duration::from_secs_f64(timeout_secs);
563
564 tokio::time::timeout(timeout, async {
565 while !self.is_active() {
566 tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
567 }
568 })
569 .await
570 .map_err(|_| {
571 KrakenWsError::ConnectionError(format!(
572 "WebSocket connection timeout after {timeout_secs} seconds"
573 ))
574 })?;
575
576 Ok(())
577 }
578
579 #[must_use]
581 pub fn is_authenticated(&self) -> bool {
582 self.auth_tracker.is_authenticated()
583 }
584
585 pub async fn wait_until_authenticated(&self, timeout_secs: f64) -> Result<(), KrakenWsError> {
589 let timeout = tokio::time::Duration::from_secs_f64(timeout_secs);
590
591 if self.auth_tracker.wait_for_authenticated(timeout).await {
592 Ok(())
593 } else {
594 Err(KrakenWsError::AuthenticationError(format!(
595 "Authentication not completed within {timeout_secs} seconds"
596 )))
597 }
598 }
599
600 pub async fn authenticate(&self) -> Result<(), KrakenWsError> {
602 if !self.config.has_api_credentials() {
603 return Err(KrakenWsError::AuthenticationError(
604 "API credentials required for authentication".to_string(),
605 ));
606 }
607
608 let _receiver = self.auth_tracker.begin();
609
610 match refresh_auth_token(&self.config).await {
611 Ok(token) => {
612 *self.auth_token.write().await = Some(token);
613 self.auth_tracker.succeed();
614 Ok(())
615 }
616 Err(e) => {
617 *self.auth_token.write().await = None;
618 self.auth_tracker.fail(e.to_string());
619 Err(e)
620 }
621 }
622 }
623
624 pub fn cancel_all_requests(&self) {
626 self.cancellation_token.cancel();
627 }
628
629 pub fn cancellation_token(&self) -> &CancellationToken {
631 &self.cancellation_token
632 }
633
634 pub async fn subscribe(
636 &self,
637 channel: KrakenWsChannel,
638 symbols: Vec<Ustr>,
639 depth: Option<u32>,
640 ) -> Result<(), KrakenWsError> {
641 if matches!(channel, KrakenWsChannel::Level3) {
642 return Err(KrakenWsError::InvalidMessage(
643 "Use subscribe_book_l3 / unsubscribe_book_l3 for the Level3 channel".to_string(),
644 ));
645 }
646 let mut symbols_to_subscribe = Vec::new();
647 let channel_str = channel.as_ref();
648 for symbol in &symbols {
649 let key = format!("{channel_str}:{symbol}");
650 if self.subscriptions.add_reference(&key) {
651 self.subscriptions.mark_subscribe(&key);
652 symbols_to_subscribe.push(*symbol);
653 }
654 }
655
656 if symbols_to_subscribe.is_empty() {
657 return Ok(());
658 }
659
660 let is_private = matches!(
661 channel,
662 KrakenWsChannel::Executions | KrakenWsChannel::Balances
663 );
664 let token = if is_private {
665 Some(self.auth_token().await.ok_or_else(|| {
666 KrakenWsError::AuthenticationError(
667 "Authentication token required for private channels. Call authenticate() first"
668 .to_string(),
669 )
670 })?)
671 } else {
672 None
673 };
674
675 let req_id = self.get_next_req_id();
676 let request = KrakenWsRequest {
677 method: KrakenWsMethod::Subscribe,
678 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
679 channel,
680 symbol: Some(symbols_to_subscribe.clone()),
681 snapshot: None,
682 depth,
683 interval: None,
684 event_trigger: None,
685 token,
686 snap_orders: None,
687 snap_trades: None,
688 })),
689 req_id: Some(req_id),
690 };
691
692 let payload = self.send_command(&request).await?;
693
694 for symbol in &symbols_to_subscribe {
695 let key = format!("{channel_str}:{symbol}");
696 self.subscriptions.confirm_subscribe(&key);
697 self.subscription_payloads
698 .write()
699 .await
700 .insert(key, payload.clone());
701 }
702
703 Ok(())
704 }
705
706 async fn subscribe_with_interval(
708 &self,
709 channel: KrakenWsChannel,
710 symbols: Vec<Ustr>,
711 interval: u32,
712 ) -> Result<(), KrakenWsError> {
713 let mut symbols_to_subscribe = Vec::new();
714 let channel_str = channel.as_ref();
715 for symbol in &symbols {
716 let key = format!("{channel_str}:{symbol}:{interval}");
717 if self.subscriptions.add_reference(&key) {
718 self.subscriptions.mark_subscribe(&key);
719 symbols_to_subscribe.push(*symbol);
720 }
721 }
722
723 if symbols_to_subscribe.is_empty() {
724 return Ok(());
725 }
726
727 let req_id = self.get_next_req_id();
728 let request = KrakenWsRequest {
729 method: KrakenWsMethod::Subscribe,
730 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
731 channel,
732 symbol: Some(symbols_to_subscribe.clone()),
733 snapshot: Some(false),
734 depth: None,
735 interval: Some(interval),
736 event_trigger: None,
737 token: None,
738 snap_orders: None,
739 snap_trades: None,
740 })),
741 req_id: Some(req_id),
742 };
743
744 let payload = self.send_command(&request).await?;
745
746 for symbol in &symbols_to_subscribe {
747 let key = format!("{channel_str}:{symbol}:{interval}");
748 self.subscriptions.confirm_subscribe(&key);
749 self.subscription_payloads
750 .write()
751 .await
752 .insert(key, payload.clone());
753 }
754
755 Ok(())
756 }
757
758 async fn unsubscribe_with_interval(
760 &self,
761 channel: KrakenWsChannel,
762 symbols: Vec<Ustr>,
763 interval: u32,
764 ) -> Result<(), KrakenWsError> {
765 let mut symbols_to_unsubscribe = Vec::new();
766 let channel_str = channel.as_ref();
767 for symbol in &symbols {
768 let key = format!("{channel_str}:{symbol}:{interval}");
769 if self.subscriptions.remove_reference(&key) {
770 self.subscriptions.mark_unsubscribe(&key);
771 symbols_to_unsubscribe.push(*symbol);
772 }
773 }
774
775 if symbols_to_unsubscribe.is_empty() {
776 return Ok(());
777 }
778
779 let req_id = self.get_next_req_id();
780 let request = KrakenWsRequest {
781 method: KrakenWsMethod::Unsubscribe,
782 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
783 channel,
784 symbol: Some(symbols_to_unsubscribe.clone()),
785 snapshot: None,
786 depth: None,
787 interval: Some(interval),
788 event_trigger: None,
789 token: None,
790 snap_orders: None,
791 snap_trades: None,
792 })),
793 req_id: Some(req_id),
794 };
795
796 self.send_command(&request).await?;
797
798 for symbol in &symbols_to_unsubscribe {
799 let key = format!("{channel_str}:{symbol}:{interval}");
800 self.subscriptions.confirm_unsubscribe(&key);
801 self.subscription_payloads.write().await.remove(&key);
802 }
803
804 Ok(())
805 }
806
807 pub async fn unsubscribe(
809 &self,
810 channel: KrakenWsChannel,
811 symbols: Vec<Ustr>,
812 ) -> Result<(), KrakenWsError> {
813 if matches!(channel, KrakenWsChannel::Level3) {
814 return Err(KrakenWsError::InvalidMessage(
815 "Use subscribe_book_l3 / unsubscribe_book_l3 for the Level3 channel".to_string(),
816 ));
817 }
818 let mut symbols_to_unsubscribe = Vec::new();
819 let channel_str = channel.as_ref();
820 for symbol in &symbols {
821 let key = format!("{channel_str}:{symbol}");
822 if self.subscriptions.remove_reference(&key) {
823 self.subscriptions.mark_unsubscribe(&key);
824 symbols_to_unsubscribe.push(*symbol);
825 } else {
826 log::debug!(
827 "Channel {channel_str} symbol {symbol} still has active subscriptions, not unsubscribing"
828 );
829 }
830 }
831
832 if symbols_to_unsubscribe.is_empty() {
833 return Ok(());
834 }
835
836 let is_private = matches!(
837 channel,
838 KrakenWsChannel::Executions | KrakenWsChannel::Balances
839 );
840 let token = if is_private {
841 Some(self.auth_token().await.ok_or_else(|| {
842 KrakenWsError::AuthenticationError(
843 "Authentication token required for private channels. Call authenticate() first"
844 .to_string(),
845 )
846 })?)
847 } else {
848 None
849 };
850
851 let req_id = self.get_next_req_id();
852 let request = KrakenWsRequest {
853 method: KrakenWsMethod::Unsubscribe,
854 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
855 channel,
856 symbol: Some(symbols_to_unsubscribe.clone()),
857 snapshot: None,
858 depth: None,
859 interval: None,
860 event_trigger: None,
861 token,
862 snap_orders: None,
863 snap_trades: None,
864 })),
865 req_id: Some(req_id),
866 };
867
868 self.send_command(&request).await?;
869
870 for symbol in &symbols_to_unsubscribe {
871 let key = format!("{channel_str}:{symbol}");
872 self.subscriptions.confirm_unsubscribe(&key);
873 self.subscription_payloads.write().await.remove(&key);
874 }
875
876 Ok(())
877 }
878
879 pub async fn send_ping(&self) -> Result<(), KrakenWsError> {
881 let req_id = self.get_next_req_id();
882
883 let request = KrakenWsRequest {
884 method: KrakenWsMethod::Ping,
885 params: None,
886 req_id: Some(req_id),
887 };
888
889 self.send_command(&request).await?;
890 Ok(())
891 }
892
893 async fn send_command(&self, request: &KrakenWsRequest) -> Result<SecretString, KrakenWsError> {
894 let payload = SecretString::from(
895 serde_json::to_string(request).map_err(|e| KrakenWsError::JsonError(e.to_string()))?,
896 );
897
898 log::trace!(
899 "Sending WebSocket request: method={:?} ({} bytes)",
900 request.method,
901 payload.expose_secret().len(),
902 );
903
904 let cmd = match request.method {
905 KrakenWsMethod::Subscribe => SpotHandlerCommand::Subscribe {
906 payload: payload.clone(),
907 },
908 KrakenWsMethod::Unsubscribe => SpotHandlerCommand::Unsubscribe {
909 payload: payload.clone(),
910 },
911 KrakenWsMethod::Ping | KrakenWsMethod::Pong => SpotHandlerCommand::Ping {
912 payload: payload.clone(),
913 },
914 KrakenWsMethod::AddOrder
915 | KrakenWsMethod::AmendOrder
916 | KrakenWsMethod::CancelOrder
917 | KrakenWsMethod::BatchAdd => {
918 return Err(KrakenWsError::InvalidMessage(
919 "Order methods must not be sent via send_command; use the dedicated order submission path".to_string()
920 ));
921 }
922 };
923
924 self.cmd_tx
925 .read()
926 .await
927 .send(cmd)
928 .map_err(|e| KrakenWsError::ConnectionError(format!("Failed to send request: {e}")))?;
929
930 Ok(payload)
931 }
932
933 pub fn is_connected(&self) -> bool {
935 let connection_mode_arc = self.connection_mode.load();
936 !ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
937 }
938
939 pub fn is_active(&self) -> bool {
941 let connection_mode_arc = self.connection_mode.load();
942 ConnectionMode::from_atomic(&connection_mode_arc).is_active()
943 && !self.signal.load(Ordering::Relaxed)
944 }
945
946 pub fn is_closed(&self) -> bool {
948 let connection_mode_arc = self.connection_mode.load();
949 ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
950 || self.signal.load(Ordering::Relaxed)
951 }
952
953 pub fn url(&self) -> &str {
955 &self.url
956 }
957
958 pub fn get_subscriptions(&self) -> Vec<String> {
960 self.subscriptions.all_topics()
961 }
962
963 pub fn subscriptions_contains(&self, topic: &str) -> bool {
965 self.subscriptions.all_topics().iter().any(|t| t == topic)
966 }
967
968 pub fn set_account_id(&self, account_id: AccountId) {
970 *self.account_id.write() = Some(account_id);
971 }
972
973 #[must_use]
975 pub fn account_id(&self) -> Option<AccountId> {
976 *self.account_id.read()
977 }
978
979 pub fn cache_instrument(&self, instrument: InstrumentAny) {
981 self.instruments.insert(instrument.id(), instrument);
982 }
983
984 pub fn account_id_shared(&self) -> &Arc<RwLock<Option<AccountId>>> {
986 &self.account_id
987 }
988
989 pub fn truncated_id_map(&self) -> &Arc<AtomicMap<String, ClientOrderId>> {
991 &self.truncated_id_map
992 }
993
994 pub fn cache_client_order(
996 &self,
997 client_order_id: ClientOrderId,
998 _venue_order_id: Option<VenueOrderId>,
999 _instrument_id: InstrumentId,
1000 _trader_id: TraderId,
1001 _strategy_id: StrategyId,
1002 ) {
1003 let truncated = crate::common::parse::truncate_cl_ord_id(&client_order_id);
1004
1005 if truncated != client_order_id.as_str() {
1006 self.truncated_id_map.insert(truncated, client_order_id);
1007 }
1008 }
1009
1010 pub fn stream(
1018 &mut self,
1019 ) -> Result<impl futures_util::Stream<Item = KrakenSpotWsMessage> + use<>, KrakenWsError> {
1020 let rx = self.out_rx.take().ok_or_else(|| {
1021 KrakenWsError::ChannelError(
1022 "Stream receiver already taken or client not connected".to_string(),
1023 )
1024 })?;
1025 let mut rx = Arc::try_unwrap(rx).map_err(|_| {
1026 KrakenWsError::ChannelError(
1027 "Cannot take ownership of stream - other client clones still hold references"
1028 .to_string(),
1029 )
1030 })?;
1031 Ok(async_stream::stream! {
1032 while let Some(msg) = rx.recv().await {
1033 yield msg;
1034 }
1035 })
1036 }
1037
1038 pub async fn subscribe_book(
1040 &self,
1041 instrument_id: InstrumentId,
1042 depth: Option<u32>,
1043 ) -> Result<(), KrakenWsError> {
1044 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1045 let depth = depth.unwrap_or(10);
1046
1047 if !matches!(depth, 10 | 25 | 100 | 500 | 1000) {
1048 return Err(KrakenWsError::InvalidMessage(format!(
1049 "Invalid L2 depth {depth}, valid values: 10, 25, 100, 500, 1000",
1050 )));
1051 }
1052
1053 let channel_str = KrakenWsChannel::Book.as_ref();
1054 let key = format!("{channel_str}:{symbol}");
1055 let is_first_reference = self.subscriptions.add_reference(&key);
1056
1057 if !is_first_reference {
1058 let existing_depth = self.l2_depths.get(symbol.as_str());
1059
1060 if existing_depth != Some(depth) {
1061 self.subscriptions.remove_reference(&key);
1062 return Err(KrakenWsError::InvalidMessage(format!(
1063 "L2 subscription for {symbol} already exists with depth \
1064 {existing_depth:?}, cannot resubscribe with depth {depth}",
1065 )));
1066 }
1067 return Ok(());
1068 }
1069
1070 self.subscriptions.mark_subscribe(&key);
1071 self.l2_depths.insert(symbol.as_str(), depth);
1072
1073 let req_id = self.get_next_req_id();
1074 let request = KrakenWsRequest {
1075 method: KrakenWsMethod::Subscribe,
1076 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1077 channel: KrakenWsChannel::Book,
1078 symbol: Some(vec![symbol]),
1079 snapshot: None,
1080 depth: Some(depth),
1081 interval: None,
1082 event_trigger: None,
1083 token: None,
1084 snap_orders: None,
1085 snap_trades: None,
1086 })),
1087 req_id: Some(req_id),
1088 };
1089
1090 let payload = match self.send_command(&request).await {
1091 Ok(payload) => payload,
1092 Err(e) => {
1093 self.l2_depths.remove(symbol.as_str());
1094 self.subscriptions.remove_reference(&key);
1095 self.subscriptions.mark_unsubscribe(&key);
1096 self.subscriptions.confirm_unsubscribe(&key);
1097 return Err(e);
1098 }
1099 };
1100
1101 self.subscriptions.confirm_subscribe(&key);
1102 self.subscription_payloads
1103 .write()
1104 .await
1105 .insert(key, payload);
1106 Ok(())
1107 }
1108
1109 pub async fn subscribe_book_l3(&self, symbol: Ustr, depth: u32) -> Result<(), KrakenWsError> {
1128 if !matches!(depth, 10 | 100 | 1000) {
1129 return Err(KrakenWsError::InvalidMessage(format!(
1130 "Invalid L3 depth {depth}, valid values: 10, 100, 1000",
1131 )));
1132 }
1133
1134 let token = self.auth_token().await.ok_or_else(|| {
1135 KrakenWsError::AuthenticationError(
1136 "Authentication token required for level3. Call authenticate() first".to_string(),
1137 )
1138 })?;
1139
1140 let channel_str = KrakenWsChannel::Level3.as_ref();
1141 let key = format!("{channel_str}:{symbol}");
1142
1143 let is_first_reference = self.subscriptions.add_reference(&key);
1144
1145 if !is_first_reference {
1146 let existing_depth = self.l3_depths.lock().get(symbol.as_str()).copied();
1147
1148 if existing_depth != Some(depth) {
1149 self.subscriptions.remove_reference(&key);
1150 return Err(KrakenWsError::InvalidMessage(format!(
1151 "L3 subscription for {symbol} already exists with depth \
1152 {existing_depth:?}, cannot resubscribe with depth {depth}",
1153 )));
1154 }
1155 return Ok(());
1156 }
1157
1158 self.subscriptions.mark_subscribe(&key);
1159
1160 self.l3_depths.lock().insert(symbol.to_string(), depth);
1161
1162 let req_id = self.get_next_req_id();
1163 let request = KrakenWsRequest {
1164 method: KrakenWsMethod::Subscribe,
1165 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1166 channel: KrakenWsChannel::Level3,
1167 symbol: Some(vec![symbol]),
1168 snapshot: Some(true),
1169 depth: Some(depth),
1170 interval: None,
1171 event_trigger: None,
1172 token: Some(token),
1173 snap_orders: None,
1174 snap_trades: None,
1175 })),
1176 req_id: Some(req_id),
1177 };
1178
1179 let payload = match self.send_command(&request).await {
1180 Ok(p) => p,
1181 Err(e) => {
1182 self.l3_depths.lock().remove(symbol.as_str());
1183 self.subscriptions.remove_reference(&key);
1184 self.subscriptions.mark_unsubscribe(&key);
1185 self.subscriptions.confirm_unsubscribe(&key);
1186 return Err(e);
1187 }
1188 };
1189
1190 self.subscriptions.confirm_subscribe(&key);
1191 self.subscription_payloads
1192 .write()
1193 .await
1194 .insert(key, payload);
1195 Ok(())
1196 }
1197
1198 pub async fn unsubscribe_book_l3(&self, symbol: Ustr) -> Result<(), KrakenWsError> {
1213 let channel_str = KrakenWsChannel::Level3.as_ref();
1214 let key = format!("{channel_str}:{symbol}");
1215 if !self.subscriptions.remove_reference(&key) {
1216 return Ok(());
1217 }
1218 self.subscriptions.mark_unsubscribe(&key);
1219
1220 let token = self.auth_token().await;
1221 let req_id = self.get_next_req_id();
1222 let request = KrakenWsRequest {
1223 method: KrakenWsMethod::Unsubscribe,
1224 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1225 channel: KrakenWsChannel::Level3,
1226 symbol: Some(vec![symbol]),
1227 snapshot: None,
1228 depth: None,
1229 interval: None,
1230 event_trigger: None,
1231 token,
1232 snap_orders: None,
1233 snap_trades: None,
1234 })),
1235 req_id: Some(req_id),
1236 };
1237
1238 self.send_command(&request).await?;
1239 self.subscriptions.confirm_unsubscribe(&key);
1240 self.subscription_payloads.write().await.remove(&key);
1241 self.l3_depths.lock().remove(symbol.as_str());
1242 Ok(())
1243 }
1244
1245 pub async fn resync_book_l3(&self, symbol: Ustr, depth: u32) -> Result<(), KrakenWsError> {
1259 let channel_str = KrakenWsChannel::Level3.as_ref();
1260 let key = format!("{channel_str}:{symbol}");
1261
1262 if !self.subscriptions_contains(&key) {
1268 log::debug!("Skipping L3 resync: subscription cancelled mid-retry, symbol={symbol}",);
1269 return Ok(());
1270 }
1271
1272 let new_token = refresh_auth_token(&self.config).await?;
1273 *self.auth_token.write().await = Some(new_token.clone());
1274
1275 if !self.subscriptions_contains(&key) {
1278 log::debug!(
1279 "Skipping L3 resync: subscription cancelled after token refresh, symbol={symbol}",
1280 );
1281 return Ok(());
1282 }
1283
1284 let unsub_req_id = self.get_next_req_id();
1285 let unsub = KrakenWsRequest {
1286 method: KrakenWsMethod::Unsubscribe,
1287 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1288 channel: KrakenWsChannel::Level3,
1289 symbol: Some(vec![symbol]),
1290 snapshot: None,
1291 depth: None,
1292 interval: None,
1293 event_trigger: None,
1294 token: Some(new_token.clone()),
1295 snap_orders: None,
1296 snap_trades: None,
1297 })),
1298 req_id: Some(unsub_req_id),
1299 };
1300 self.send_command(&unsub).await?;
1301
1302 if !self.subscriptions_contains(&key) {
1304 log::debug!("Skipping L3 resync resubscribe: cancelled before send, symbol={symbol}",);
1305 return Ok(());
1306 }
1307
1308 let sub_req_id = self.get_next_req_id();
1309 let sub = KrakenWsRequest {
1310 method: KrakenWsMethod::Subscribe,
1311 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1312 channel: KrakenWsChannel::Level3,
1313 symbol: Some(vec![symbol]),
1314 snapshot: Some(true),
1315 depth: Some(depth),
1316 interval: None,
1317 event_trigger: None,
1318 token: Some(new_token.clone()),
1319 snap_orders: None,
1320 snap_trades: None,
1321 })),
1322 req_id: Some(sub_req_id),
1323 };
1324 let payload = self.send_command(&sub).await?;
1325
1326 if self.subscriptions_contains(&key) {
1331 self.subscription_payloads
1332 .write()
1333 .await
1334 .insert(key, payload);
1335 self.l3_depths.lock().insert(symbol.to_string(), depth);
1336 }
1337
1338 Ok(())
1339 }
1340
1341 pub fn validate_l3_checksum(&self) -> bool {
1343 self.config.validate_l3_checksum
1344 }
1345
1346 pub fn has_credentials(&self) -> bool {
1349 self.config.has_api_credentials()
1350 }
1351
1352 pub fn instruments_handle(&self) -> Arc<AtomicMap<InstrumentId, InstrumentAny>> {
1358 Arc::clone(&self.instruments)
1359 }
1360
1361 pub fn l3_depths_handle(&self) -> Arc<parking_lot::Mutex<ahash::AHashMap<String, u32>>> {
1366 Arc::clone(&self.l3_depths)
1367 }
1368
1369 pub(crate) fn l2_depths_handle(&self) -> L2Depths {
1371 self.l2_depths.clone()
1372 }
1373
1374 pub async fn subscribe_quotes(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1379 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1380 let key = format!("quotes:{symbol}");
1381
1382 if !self.subscriptions.add_reference(&key) {
1383 return Ok(());
1384 }
1385
1386 self.subscriptions.mark_subscribe(&key);
1387
1388 let req_id = self.get_next_req_id();
1389 let request = KrakenWsRequest {
1390 method: KrakenWsMethod::Subscribe,
1391 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1392 channel: KrakenWsChannel::Ticker,
1393 symbol: Some(vec![symbol]),
1394 snapshot: None,
1395 depth: None,
1396 interval: None,
1397 event_trigger: Some("bbo".to_string()),
1398 token: None,
1399 snap_orders: None,
1400 snap_trades: None,
1401 })),
1402 req_id: Some(req_id),
1403 };
1404
1405 let payload = self.send_command(&request).await?;
1406 self.subscriptions.confirm_subscribe(&key);
1407 self.subscription_payloads
1408 .write()
1409 .await
1410 .insert(key, payload);
1411 Ok(())
1412 }
1413
1414 pub async fn subscribe_trades(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1416 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1417 self.subscribe(KrakenWsChannel::Trade, vec![symbol], None)
1418 .await
1419 }
1420
1421 pub async fn subscribe_bars(&self, bar_type: BarType) -> Result<(), KrakenWsError> {
1427 let symbol = to_ws_v2_symbol(bar_type.instrument_id().symbol.inner());
1428 let interval = bar_type_to_ws_interval(bar_type)?;
1429 self.subscribe_with_interval(KrakenWsChannel::Ohlc, vec![symbol], interval)
1430 .await
1431 }
1432
1433 pub async fn subscribe_executions(
1437 &self,
1438 snap_orders: bool,
1439 snap_trades: bool,
1440 ) -> Result<(), KrakenWsError> {
1441 let req_id = self.get_next_req_id();
1442
1443 let token = self.auth_token().await.ok_or_else(|| {
1444 KrakenWsError::AuthenticationError(
1445 "Authentication token required for executions channel. Call authenticate() first"
1446 .to_string(),
1447 )
1448 })?;
1449
1450 let request = KrakenWsRequest {
1451 method: KrakenWsMethod::Subscribe,
1452 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1453 channel: KrakenWsChannel::Executions,
1454 symbol: None,
1455 snapshot: None,
1456 depth: None,
1457 interval: None,
1458 event_trigger: None,
1459 token: Some(token),
1460 snap_orders: Some(snap_orders),
1461 snap_trades: Some(snap_trades),
1462 })),
1463 req_id: Some(req_id),
1464 };
1465
1466 let payload = self.send_command(&request).await?;
1467
1468 let key = "executions";
1469 if self.subscriptions.add_reference(key) {
1470 self.subscriptions.mark_subscribe(key);
1471 self.subscriptions.confirm_subscribe(key);
1472 self.subscription_payloads
1473 .write()
1474 .await
1475 .insert(key.to_string(), payload);
1476 }
1477
1478 Ok(())
1479 }
1480
1481 pub async fn unsubscribe_book(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1483 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1484 let channel_str = KrakenWsChannel::Book.as_ref();
1485 let key = format!("{channel_str}:{symbol}");
1486
1487 if !self.subscriptions.remove_reference(&key) {
1488 return Ok(());
1489 }
1490
1491 self.subscriptions.mark_unsubscribe(&key);
1492
1493 let req_id = self.get_next_req_id();
1494 let request = KrakenWsRequest {
1495 method: KrakenWsMethod::Unsubscribe,
1496 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1497 channel: KrakenWsChannel::Book,
1498 symbol: Some(vec![symbol]),
1499 snapshot: None,
1500 depth: None,
1501 interval: None,
1502 event_trigger: None,
1503 token: None,
1504 snap_orders: None,
1505 snap_trades: None,
1506 })),
1507 req_id: Some(req_id),
1508 };
1509
1510 self.send_command(&request).await?;
1511 self.subscriptions.confirm_unsubscribe(&key);
1512 self.subscription_payloads.write().await.remove(&key);
1513 self.l2_depths.remove(symbol.as_str());
1514 Ok(())
1515 }
1516
1517 pub async fn unsubscribe_quotes(
1519 &self,
1520 instrument_id: InstrumentId,
1521 ) -> Result<(), KrakenWsError> {
1522 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1523 let key = format!("quotes:{symbol}");
1524
1525 if !self.subscriptions.remove_reference(&key) {
1526 return Ok(());
1527 }
1528
1529 self.subscriptions.mark_unsubscribe(&key);
1530
1531 let req_id = self.get_next_req_id();
1532 let request = KrakenWsRequest {
1533 method: KrakenWsMethod::Unsubscribe,
1534 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1535 channel: KrakenWsChannel::Ticker,
1536 symbol: Some(vec![symbol]),
1537 snapshot: None,
1538 depth: None,
1539 interval: None,
1540 event_trigger: Some("bbo".to_string()),
1541 token: None,
1542 snap_orders: None,
1543 snap_trades: None,
1544 })),
1545 req_id: Some(req_id),
1546 };
1547
1548 self.send_command(&request).await?;
1549 self.subscriptions.confirm_unsubscribe(&key);
1550 self.subscription_payloads.write().await.remove(&key);
1551 Ok(())
1552 }
1553
1554 pub async fn unsubscribe_trades(
1556 &self,
1557 instrument_id: InstrumentId,
1558 ) -> Result<(), KrakenWsError> {
1559 let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1560 self.unsubscribe(KrakenWsChannel::Trade, vec![symbol]).await
1561 }
1562
1563 pub async fn unsubscribe_bars(&self, bar_type: BarType) -> Result<(), KrakenWsError> {
1569 let symbol = to_ws_v2_symbol(bar_type.instrument_id().symbol.inner());
1570 let interval = bar_type_to_ws_interval(bar_type)?;
1571 self.unsubscribe_with_interval(KrakenWsChannel::Ohlc, vec![symbol], interval)
1572 .await
1573 }
1574}
1575
1576async fn refresh_auth_token(
1578 config: &KrakenDataClientConfig,
1579) -> Result<SecretString, KrakenWsError> {
1580 let api_key = config
1581 .api_key
1582 .clone()
1583 .ok_or_else(|| KrakenWsError::AuthenticationError("Missing API key".to_string()))?;
1584 let api_secret = config
1585 .api_secret
1586 .clone()
1587 .ok_or_else(|| KrakenWsError::AuthenticationError("Missing API secret".to_string()))?;
1588
1589 let http_client = KrakenSpotHttpClient::with_credentials(
1590 api_key.into_inner(),
1591 api_secret.into_inner(),
1592 config.environment,
1593 Some(config.http_base_url()),
1594 config.timeout_secs,
1595 None,
1596 None,
1597 None,
1598 config
1599 .proxy_url
1600 .as_ref()
1601 .map(|value| value.expose_secret().to_owned()),
1602 config
1603 .max_requests_per_second
1604 .unwrap_or(KRAKEN_SPOT_DEFAULT_RATE_LIMIT_PER_SECOND),
1605 )
1606 .map_err(|e| {
1607 KrakenWsError::AuthenticationError(format!("Failed to create HTTP client: {e}"))
1608 })?;
1609
1610 let ws_token = http_client.get_websockets_token().await.map_err(|e| {
1611 KrakenWsError::AuthenticationError(format!("Failed to get WebSocket token: {e}"))
1612 })?;
1613
1614 log::debug!(
1615 "WebSocket authentication token refreshed: token_length={}, expires={}",
1616 ws_token.token.expose_secret().len(),
1617 ws_token.expires
1618 );
1619
1620 Ok(ws_token.into_token())
1621}
1622
1623fn update_auth_token_in_payload(
1624 payload: &str,
1625 new_token: &str,
1626) -> Result<SecretString, KrakenWsError> {
1627 let mut value: serde_json::Value =
1628 serde_json::from_str(payload).map_err(|e| KrakenWsError::JsonError(e.to_string()))?;
1629
1630 if let Some(params) = value.get_mut("params") {
1631 params["token"] = serde_json::Value::String(new_token.to_string());
1632 }
1633
1634 serde_json::to_string(&value)
1635 .map(SecretString::from)
1636 .map_err(|e| KrakenWsError::JsonError(e.to_string()))
1637}
1638
1639#[inline]
1640fn to_ws_v2_symbol(symbol: Ustr) -> Ustr {
1641 Ustr::from(&normalize_spot_symbol(symbol.as_str()))
1642}
1643
1644fn bar_type_to_ws_interval(bar_type: BarType) -> Result<u32, KrakenWsError> {
1645 const VALID_INTERVALS: [u32; 9] = [1, 5, 15, 30, 60, 240, 1440, 10080, 21600];
1646
1647 let spec = bar_type.spec();
1648 let step = spec.step.get() as u32;
1649
1650 let base_minutes = match spec.aggregation {
1651 BarAggregation::Minute => 1,
1652 BarAggregation::Hour => 60,
1653 BarAggregation::Day => 1440,
1654 BarAggregation::Week => 10080,
1655 other => {
1656 return Err(KrakenWsError::SubscriptionError(format!(
1657 "Unsupported bar aggregation for Kraken OHLC streaming: {other:?}"
1658 )));
1659 }
1660 };
1661
1662 let interval = base_minutes * step;
1663
1664 if !VALID_INTERVALS.contains(&interval) {
1665 return Err(KrakenWsError::SubscriptionError(format!(
1666 "Invalid bar interval {interval} minutes for Kraken OHLC streaming. \
1667 Supported intervals: 1, 5, 15, 30, 60, 240, 1440, 10080, 21600"
1668 )));
1669 }
1670
1671 Ok(interval)
1672}
1673
1674#[cfg(test)]
1675mod tests {
1676 use std::sync::{Arc, atomic::Ordering};
1677
1678 use log::{Level, LevelFilter, Log, Metadata, Record};
1679 use parking_lot::Mutex;
1680 use rstest::rstest;
1681 use tokio_util::sync::CancellationToken;
1682
1683 use super::*;
1684 use crate::config::KrakenDataClientConfig;
1685
1686 #[rstest]
1687 #[tokio::test]
1688 async fn test_debug_redacts_auth_state_and_proxy_url() {
1689 let client = KrakenSpotWebSocketClient::new(
1690 KrakenDataClientConfig::default(),
1691 CancellationToken::new(),
1692 Some("http://user:proxy-secret@localhost".to_string()),
1693 );
1694 *client.auth_token.write().await = Some(SecretString::from("auth-token".to_string()));
1695
1696 let debug = format!("{client:?}");
1697
1698 assert!(debug.contains(REDACTED));
1699 assert!(!debug.contains("proxy-secret"));
1700 assert!(!debug.contains("auth-token"));
1701 }
1702
1703 const SECRET_MARKER: &str = "OUTBOUND_SECRET_MARKER";
1704
1705 struct OutboundLogCapture {
1706 messages: Mutex<Vec<String>>,
1707 }
1708
1709 static OUTBOUND_LOG_CAPTURE: OutboundLogCapture = OutboundLogCapture {
1710 messages: Mutex::new(Vec::new()),
1711 };
1712
1713 impl OutboundLogCapture {
1714 fn clear(&self) {
1715 self.messages.lock().clear();
1716 }
1717
1718 fn messages(&self) -> Vec<String> {
1719 self.messages.lock().clone()
1720 }
1721 }
1722
1723 impl Log for OutboundLogCapture {
1724 fn enabled(&self, metadata: &Metadata<'_>) -> bool {
1725 metadata.level() == Level::Trace
1726 && metadata.target() == "nautilus_kraken::websocket::spot_v2::client"
1727 }
1728
1729 fn log(&self, record: &Record<'_>) {
1730 if self.enabled(record.metadata()) {
1731 let message = record.args().to_string();
1732 if message.starts_with("Sending WebSocket request") {
1733 self.messages.lock().push(message);
1734 }
1735 }
1736 }
1737
1738 fn flush(&self) {}
1739 }
1740
1741 #[rstest]
1742 fn test_req_id_counter_is_shared_arc_and_monotonic() {
1743 let cfg = KrakenDataClientConfig::default();
1744 let client = KrakenSpotWebSocketClient::new(cfg, CancellationToken::new(), None);
1745 let counter = client.req_id_counter();
1746 let a = counter.fetch_add(1, Ordering::Relaxed);
1747 let b = counter.fetch_add(1, Ordering::Relaxed);
1748 assert!(b > a);
1749 #[allow(clippy::redundant_clone)]
1750 let cloned = client.clone();
1751 let cloned_counter = cloned.req_id_counter();
1752 assert!(Arc::ptr_eq(&counter, &cloned_counter));
1753 }
1754
1755 #[rstest]
1756 #[case("XBT/EUR", "BTC/EUR")]
1757 #[case("XBT/USD", "BTC/USD")]
1758 #[case("XBT/USDT", "BTC/USDT")]
1759 #[case("ETH/USD", "ETH/USD")]
1760 #[case("ETH/XBT", "ETH/BTC")]
1761 #[case("SOL/XBT", "SOL/BTC")]
1762 #[case("SOL/USD", "SOL/USD")]
1763 #[case("BTC/USD", "BTC/USD")]
1764 #[case("ETH/BTC", "ETH/BTC")]
1765 #[case("XDG/USD", "DOGE/USD")]
1766 #[case("XDG/EUR", "DOGE/EUR")]
1767 fn test_to_kraken_ws_v2_symbol(#[case] input: &str, #[case] expected: &str) {
1768 let symbol = Ustr::from(input);
1769 let result = to_ws_v2_symbol(symbol);
1770 assert_eq!(result, expected);
1771 }
1772
1773 fn test_client_without_credentials() -> KrakenSpotWebSocketClient {
1774 KrakenSpotWebSocketClient::new(
1775 KrakenDataClientConfig::default(),
1776 CancellationToken::new(),
1777 None,
1778 )
1779 }
1780
1781 #[rstest]
1782 #[tokio::test]
1783 async fn test_outbound_logs_omit_payload_bodies() {
1784 log::set_logger(&OUTBOUND_LOG_CAPTURE).expect("test logger already installed");
1785 log::set_max_level(LevelFilter::Trace);
1786
1787 let client = test_client_without_credentials();
1788 let request = KrakenWsRequest {
1789 method: KrakenWsMethod::Subscribe,
1790 params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1791 channel: KrakenWsChannel::Executions,
1792 symbol: None,
1793 snapshot: None,
1794 depth: None,
1795 interval: None,
1796 event_trigger: None,
1797 token: Some(SecretString::from(SECRET_MARKER)),
1798 snap_orders: Some(true),
1799 snap_trades: Some(false),
1800 })),
1801 req_id: Some(426),
1802 };
1803 let payload_len = serde_json::to_string(&request).unwrap().len();
1804 OUTBOUND_LOG_CAPTURE.clear();
1805
1806 let error = client.send_command(&request).await.unwrap_err();
1807 let messages = OUTBOUND_LOG_CAPTURE.messages();
1808
1809 assert!(matches!(error, KrakenWsError::ConnectionError(_)));
1810 assert!(
1811 messages
1812 .iter()
1813 .all(|message| !message.contains(SECRET_MARKER)),
1814 "outbound logs exposed the secret marker: {messages:?}"
1815 );
1816 assert!(
1817 messages.iter().any(|message| {
1818 message
1819 == &format!("Sending WebSocket request: method=Subscribe ({payload_len} bytes)")
1820 }),
1821 "subscribe metadata missing or inaccurate: {messages:?}"
1822 );
1823 }
1824
1825 #[rstest]
1826 #[tokio::test]
1827 async fn test_authenticate_without_credentials_errors() {
1828 let client = test_client_without_credentials();
1829
1830 let err = client.authenticate().await.expect_err("should fail");
1831 assert!(
1832 matches!(err, KrakenWsError::AuthenticationError(ref msg) if msg.contains("API credentials required")),
1833 "unexpected error: {err:?}"
1834 );
1835 assert!(!client.is_authenticated());
1836 }
1837
1838 #[rstest]
1839 #[tokio::test]
1840 async fn test_wait_until_authenticated_times_out() {
1841 let client = test_client_without_credentials();
1842
1843 let err = client
1844 .wait_until_authenticated(0.05)
1845 .await
1846 .expect_err("should time out");
1847 assert!(matches!(err, KrakenWsError::AuthenticationError(_)));
1848 }
1849
1850 #[rstest]
1851 #[tokio::test]
1852 async fn test_wait_until_authenticated_resolves_after_succeed() {
1853 let client = test_client_without_credentials();
1854
1855 let tracker = client.auth_tracker.clone();
1856 let _rx = tracker.begin();
1857
1858 tokio::spawn(async move {
1859 tokio::time::sleep(tokio::time::Duration::from_millis(20)).await;
1860 tracker.succeed();
1861 });
1862
1863 client
1864 .wait_until_authenticated(1.0)
1865 .await
1866 .expect("should resolve once tracker succeeds");
1867 assert!(client.is_authenticated());
1868 }
1869
1870 #[rstest]
1871 #[tokio::test]
1872 async fn test_is_authenticated_flips_on_fail() {
1873 let client = test_client_without_credentials();
1874
1875 let _rx = client.auth_tracker.begin();
1876 client.auth_tracker.succeed();
1877 assert!(client.is_authenticated());
1878
1879 client.auth_tracker.fail("test failure");
1880 assert!(!client.is_authenticated());
1881 }
1882
1883 #[rstest]
1884 fn test_l3_factory_uses_ws_l3_url() {
1885 let cfg = KrakenDataClientConfig::default();
1886 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1887 assert_eq!(client.url(), "wss://ws-l3.kraken.com/v2");
1888 }
1889
1890 #[rstest]
1891 fn test_l3_factory_respects_override() {
1892 let cfg = KrakenDataClientConfig {
1893 ws_l3_url: Some("wss://override.example/v2".to_string()),
1894 ..Default::default()
1895 };
1896 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1897 assert_eq!(client.url(), "wss://override.example/v2");
1898 }
1899
1900 #[rstest]
1901 #[tokio::test]
1902 async fn test_subscribe_book_l3_without_auth_errors_and_leaves_clean_state() {
1903 let cfg = KrakenDataClientConfig::default();
1904 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1905
1906 let err = client
1907 .subscribe_book_l3(Ustr::from("BTC/USD"), 1000)
1908 .await
1909 .expect_err("should fail without auth token");
1910
1911 assert!(matches!(err, KrakenWsError::AuthenticationError(_)));
1912 assert!(
1913 client.subscriptions.is_empty(),
1914 "no state must leak on auth failure"
1915 );
1916 }
1917
1918 #[rstest]
1919 #[tokio::test]
1920 async fn test_subscribe_book_l3_invalid_depth_errors() {
1921 let cfg = KrakenDataClientConfig::default();
1922 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
1923
1924 let err = client
1925 .subscribe_book_l3(Ustr::from("BTC/USD"), 50)
1926 .await
1927 .expect_err("should fail on invalid depth");
1928
1929 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
1930 assert!(!client.subscriptions_contains("level3:BTC/USD"));
1931 }
1932
1933 #[rstest]
1934 #[tokio::test]
1935 async fn test_subscribe_book_invalid_depth_errors() {
1936 let client = test_client_without_credentials();
1937 let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
1938
1939 let err = client
1940 .subscribe_book(instrument_id, Some(50))
1941 .await
1942 .expect_err("should fail on invalid depth");
1943
1944 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
1945 assert!(!client.subscriptions_contains("book:BTC/USD"));
1946 assert_eq!(client.l2_depths.get("BTC/USD"), None);
1947 }
1948
1949 #[rstest]
1950 #[tokio::test]
1951 async fn test_subscribe_book_defaults_to_depth_10_and_stores_state() {
1952 let client = test_client_without_credentials();
1953 let key = "book:BTC/USD";
1954 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1955 *client.cmd_tx.write().await = cmd_tx;
1956
1957 client
1958 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), None)
1959 .await
1960 .expect("subscribe should succeed");
1961
1962 let cmd = cmd_rx.try_recv().expect("expected subscribe command");
1963 let SpotHandlerCommand::Subscribe { payload } = cmd else {
1964 panic!("expected subscribe command");
1965 };
1966 assert_book_subscribe_payload(payload.expose_secret(), "BTC/USD", 10);
1967
1968 assert_eq!(client.subscriptions.get_reference_count(key), 1);
1969 assert!(client.subscriptions_contains(key));
1970 assert_eq!(client.l2_depths.get("BTC/USD"), Some(10));
1971 assert_eq!(
1972 client.subscription_payloads.read().await.get(key),
1973 Some(&payload)
1974 );
1975 }
1976
1977 #[rstest]
1978 #[tokio::test]
1979 async fn test_subscribe_book_send_failure_rolls_back_l2_state() {
1980 let client = test_client_without_credentials();
1981 let key = "book:BTC/USD";
1982
1983 let err = client
1984 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(10))
1985 .await
1986 .expect_err("should fail when command receiver is closed");
1987
1988 assert!(matches!(err, KrakenWsError::ConnectionError(_)));
1989 assert_eq!(client.subscriptions.get_reference_count(key), 0);
1990 assert!(!client.subscriptions_contains(key));
1991 assert!(client.subscriptions.pending_subscribe_topics().is_empty());
1992 assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
1993 assert_eq!(client.l2_depths.get("BTC/USD"), None);
1994 assert!(!client.subscription_payloads.read().await.contains_key(key));
1995 }
1996
1997 #[rstest]
1998 fn test_subscribe_book_l3_refcount_idempotent() {
1999 let cfg = KrakenDataClientConfig::default();
2000 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2001 let key = "level3:BTC/USD";
2002
2003 assert!(client.subscriptions.add_reference(key));
2004 assert!(!client.subscriptions.add_reference(key));
2005 assert!(!client.subscriptions.remove_reference(key));
2006 assert!(client.subscriptions.remove_reference(key));
2007 }
2008
2009 #[rstest]
2010 #[tokio::test]
2011 async fn test_subscribe_book_l3_rejects_depth_mismatch() {
2012 let cfg = KrakenDataClientConfig::default();
2013 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2014
2015 let key = "level3:BTC/USD";
2016 client.subscriptions.add_reference(key);
2017 client.subscriptions.mark_subscribe(key);
2018 client.subscriptions.confirm_subscribe(key);
2019 client.l3_depths.lock().insert("BTC/USD".to_string(), 1000);
2020
2021 *client.auth_token.write().await = Some(SecretString::from("test-token".to_string()));
2022
2023 let err = client
2024 .subscribe_book_l3(Ustr::from("BTC/USD"), 10)
2025 .await
2026 .expect_err("should reject depth mismatch");
2027 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2028
2029 assert!(client.subscriptions.remove_reference(key));
2030 }
2031
2032 #[rstest]
2033 #[tokio::test]
2034 async fn test_subscribe_book_rejects_depth_mismatch() {
2035 let client = test_client_without_credentials();
2036 let key = "book:BTC/USD";
2037 client.subscriptions.add_reference(key);
2038 client.subscriptions.mark_subscribe(key);
2039 client.subscriptions.confirm_subscribe(key);
2040 client.l2_depths.insert("BTC/USD", 10);
2041
2042 let err = client
2043 .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(25))
2044 .await
2045 .expect_err("should reject depth mismatch");
2046 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2047
2048 assert!(client.subscriptions.remove_reference(key));
2049 }
2050
2051 #[rstest]
2052 #[tokio::test]
2053 async fn test_unsubscribe_book_removes_l2_depth_on_last_reference() {
2054 let client = test_client_without_credentials();
2055 let key = "book:BTC/USD";
2056 let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2057 *client.cmd_tx.write().await = cmd_tx;
2058
2059 client.subscriptions.add_reference(key);
2060 client.subscriptions.mark_subscribe(key);
2061 client.subscriptions.confirm_subscribe(key);
2062 client.l2_depths.insert("BTC/USD", 10);
2063 client
2064 .subscription_payloads
2065 .write()
2066 .await
2067 .insert(key.to_string(), SecretString::from("payload".to_string()));
2068
2069 client
2070 .unsubscribe_book(InstrumentId::from("BTC/USD.KRAKEN"))
2071 .await
2072 .expect("unsubscribe should succeed");
2073
2074 assert_eq!(client.subscriptions.get_reference_count(key), 0);
2075 assert!(!client.subscriptions_contains(key));
2076 assert_eq!(client.l2_depths.get("BTC/USD"), None);
2077 assert!(!client.subscription_payloads.read().await.contains_key(key));
2078
2079 let cmd = cmd_rx.try_recv().expect("expected unsubscribe command");
2080 let SpotHandlerCommand::Unsubscribe { payload } = cmd else {
2081 panic!("expected unsubscribe command");
2082 };
2083 assert!(
2084 payload
2085 .expose_secret()
2086 .contains(r#""method":"unsubscribe""#),
2087 );
2088 assert!(payload.expose_secret().contains(r#""channel":"book""#));
2089 assert!(payload.expose_secret().contains(r#""BTC/USD""#));
2090 }
2091
2092 #[rstest]
2093 #[tokio::test]
2094 async fn test_generic_subscribe_rejects_level3() {
2095 let cfg = KrakenDataClientConfig::default();
2096 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2097
2098 let err = client
2099 .subscribe(
2100 KrakenWsChannel::Level3,
2101 vec![Ustr::from("BTC/USD")],
2102 Some(1000),
2103 )
2104 .await
2105 .expect_err("generic subscribe must reject Level3");
2106 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2107
2108 let err = client
2109 .unsubscribe(KrakenWsChannel::Level3, vec![Ustr::from("BTC/USD")])
2110 .await
2111 .expect_err("generic unsubscribe must reject Level3");
2112 assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2113 }
2114
2115 #[rstest]
2116 fn test_update_auth_token_in_payload_for_level3() {
2117 let original = r#"{"method":"subscribe","params":{"channel":"level3","symbol":["BTC/USD"],"depth":1000,"snapshot":true,"token":"OLD"},"req_id":1}"#;
2118 let rewritten = update_auth_token_in_payload(original, "NEW").unwrap();
2119 assert!(rewritten.expose_secret().contains(r#""token":"NEW""#));
2120 assert!(!rewritten.expose_secret().contains(r#""token":"OLD""#));
2121 }
2122
2123 #[rstest]
2124 fn test_l3_depths_shared_between_subscribe_and_handle() {
2125 let cfg = KrakenDataClientConfig::default();
2126 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2127
2128 let handle = client.l3_depths_handle();
2129 client.l3_depths.lock().insert("BTC/USD".to_string(), 100);
2130
2131 assert_eq!(handle.lock().get("BTC/USD").copied(), Some(100));
2132 }
2133
2134 #[rstest]
2135 fn test_l2_depths_shared_between_subscribe_and_handle() {
2136 let client = test_client_without_credentials();
2137
2138 let handle = client.l2_depths_handle();
2139 client.l2_depths.insert("BTC/USD", 10);
2140
2141 assert_eq!(handle.get("BTC/USD"), Some(10));
2142 }
2143
2144 #[rstest]
2145 fn test_resync_book_l3_does_not_touch_refcount() {
2146 let cfg = KrakenDataClientConfig::default();
2147 let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2148 let key = "level3:BTC/USD";
2149
2150 assert!(client.subscriptions.add_reference(key));
2151 assert!(!client.subscriptions.add_reference(key));
2152 client.subscriptions.mark_subscribe(key);
2153 client.subscriptions.confirm_subscribe(key);
2154
2155 assert!(client.subscriptions_contains(key));
2156 }
2157
2158 fn assert_book_subscribe_payload(payload: &str, symbol: &str, depth: u32) {
2159 let value: serde_json::Value =
2160 serde_json::from_str(payload).expect("payload should parse as JSON");
2161
2162 assert_eq!(value["method"], serde_json::json!("subscribe"));
2163 assert_eq!(value["params"]["channel"], serde_json::json!("book"));
2164 assert_eq!(value["params"]["symbol"], serde_json::json!([symbol]));
2165 assert_eq!(value["params"]["depth"], serde_json::json!(depth));
2166 }
2167}