Skip to main content

nautilus_kraken/websocket/spot_v2/
client.rs

1// -------------------------------------------------------------------------------------------------
2//  Copyright (C) 2015-2026 Nautech Systems Pty Ltd. All rights reserved.
3//  https://nautechsystems.io
4//
5//  Licensed under the GNU Lesser General Public License Version 3.0 (the "License");
6//  You may not use this file except in compliance with the License.
7//  You may obtain a copy of the License at https://www.gnu.org/licenses/lgpl-3.0.en.html
8//
9//  Unless required by applicable law or agreed to in writing, software
10//  distributed under the License is distributed on an "AS IS" BASIS,
11//  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12//  See the License for the specific language governing permissions and
13//  limitations under the License.
14// -------------------------------------------------------------------------------------------------
15
16//! WebSocket client for the Kraken v2 streaming API.
17
18use 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    ConnectionUrl,
42    http::create_standard_nautilus_headers,
43    mode::ConnectionMode,
44    websocket::{
45        AuthTracker, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
46        channel_message_handler,
47    },
48};
49use parking_lot::RwLock;
50use tokio_util::sync::CancellationToken;
51use ustr::Ustr;
52
53/// Topic delimiter for Kraken Spot v2 WebSocket subscriptions.
54///
55/// Topics use colon format: `channel:symbol` (e.g., `Trade:ETH/USD`).
56pub const KRAKEN_SPOT_WS_TOPIC_DELIMITER: char = ':';
57
58use super::{
59    enums::{KrakenWsChannel, KrakenWsMethod},
60    handler::{SpotFeedHandler, SpotHandlerCommand},
61    level_2::{L2BookRequest, L2BookRequests, L2Depths},
62    messages::{KrakenSpotWsMessage, KrakenWsChannelParams, KrakenWsParams, KrakenWsRequest},
63};
64use crate::{
65    common::{
66        consts::{
67            KRAKEN_RATE_LIMIT_KEY_ORDER, KRAKEN_RATE_LIMIT_KEY_SUBSCRIPTION,
68            KRAKEN_SPOT_WS_ORDER_QUOTA, KRAKEN_SPOT_WS_SUBSCRIPTION_QUOTA,
69        },
70        parse::normalize_spot_symbol,
71    },
72    config::KrakenDataClientConfig,
73    http::{KrakenSpotHttpClient, spot::client::KRAKEN_SPOT_DEFAULT_RATE_LIMIT_PER_SECOND},
74    websocket::error::KrakenWsError,
75};
76
77const WS_PING_MSG: &str = r#"{"method":"ping"}"#;
78
79/// WebSocket client for the Kraken Spot v2 streaming API.
80#[derive(Debug)]
81pub struct KrakenSpotWebSocketClient {
82    url: ConnectionUrl,
83    config: KrakenDataClientConfig,
84    signal: Arc<AtomicBool>,
85    connection_mode: Arc<ArcSwap<AtomicU8>>,
86    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand>>>,
87    out_rx: Option<Arc<tokio::sync::mpsc::UnboundedReceiver<KrakenSpotWsMessage>>>,
88    handler_tasks: Arc<TaskGroup>,
89    connect_lock: Arc<tokio::sync::Mutex<()>>,
90    subscriptions: SubscriptionState,
91    subscription_payloads: Arc<tokio::sync::RwLock<HashMap<String, SecretString>>>,
92    auth_tracker: AuthTracker,
93    cancellation_token: CancellationToken,
94    req_id_counter: Arc<AtomicU64>,
95    auth_token: Arc<tokio::sync::RwLock<Option<SecretString>>>,
96    account_id: Arc<RwLock<Option<AccountId>>>,
97    truncated_id_map: Arc<AtomicMap<String, ClientOrderId>>,
98    instruments: Arc<AtomicMap<InstrumentId, InstrumentAny>>,
99    l2_depths: L2Depths,
100    /// Serializes `subscribe_book`, `unsubscribe_book` and `resync_book`, so a recovery's check
101    /// of the live subscription and its two commands cannot interleave with a replacement.
102    l2_book_commands: Arc<tokio::sync::Mutex<()>>,
103    /// Every `book` subscribe sent and not yet answered, by request id, so the data client can
104    /// match the venue's answer to the subscription behind it.
105    book_requests: L2BookRequests,
106    l3_depths: Arc<parking_lot::Mutex<ahash::AHashMap<String, u32>>>,
107    transport_backend: TransportBackend,
108    proxy_url: Option<SecretString>,
109    socket_control: Option<SocketControl>,
110}
111
112impl Clone for KrakenSpotWebSocketClient {
113    fn clone(&self) -> Self {
114        Self {
115            url: self.url.clone(),
116            config: self.config.clone(),
117            signal: Arc::clone(&self.signal),
118            connection_mode: Arc::clone(&self.connection_mode),
119            cmd_tx: Arc::clone(&self.cmd_tx),
120            out_rx: self.out_rx.clone(),
121            handler_tasks: Arc::clone(&self.handler_tasks),
122            connect_lock: Arc::clone(&self.connect_lock),
123            subscriptions: self.subscriptions.clone(),
124            subscription_payloads: Arc::clone(&self.subscription_payloads),
125            auth_tracker: self.auth_tracker.clone(),
126            cancellation_token: self.cancellation_token.clone(),
127            req_id_counter: self.req_id_counter.clone(),
128            auth_token: self.auth_token.clone(),
129            account_id: Arc::clone(&self.account_id),
130            truncated_id_map: Arc::clone(&self.truncated_id_map),
131            instruments: Arc::clone(&self.instruments),
132            l2_depths: self.l2_depths.clone(),
133            l2_book_commands: Arc::clone(&self.l2_book_commands),
134            book_requests: Arc::clone(&self.book_requests),
135            l3_depths: Arc::clone(&self.l3_depths),
136            transport_backend: self.transport_backend,
137            proxy_url: self.proxy_url.clone(),
138            socket_control: self.socket_control.clone(),
139        }
140    }
141}
142
143impl KrakenSpotWebSocketClient {
144    /// Creates a new client for the configured public/private endpoint.
145    pub fn new(
146        config: KrakenDataClientConfig,
147        cancellation_token: CancellationToken,
148        proxy_url: Option<String>,
149    ) -> Self {
150        let url = if config.ws_private_url.is_some() {
151            config.ws_private_url()
152        } else {
153            config.ws_public_url()
154        };
155        Self::new_with_url(url, config, cancellation_token, proxy_url)
156    }
157
158    /// Creates a new client configured for the Kraken Spot `level3` WebSocket endpoint.
159    ///
160    /// Selects `config.ws_l3_url()` and otherwise mirrors [`Self::new`]. `Level3`
161    /// subscriptions are treated as authenticated and must follow `authenticate()`.
162    pub fn l3(
163        config: KrakenDataClientConfig,
164        cancellation_token: CancellationToken,
165        proxy_url: Option<String>,
166    ) -> Self {
167        let url = config.ws_l3_url();
168        Self::new_with_url(url, config, cancellation_token, proxy_url)
169    }
170
171    fn new_with_url(
172        url: String,
173        mut config: KrakenDataClientConfig,
174        cancellation_token: CancellationToken,
175        proxy_url: Option<String>,
176    ) -> Self {
177        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel::<SpotHandlerCommand>();
178        let initial_mode = AtomicU8::new(ConnectionMode::Closed.as_u8());
179        let connection_mode = Arc::new(ArcSwap::from_pointee(initial_mode));
180
181        let transport_backend = config.transport_backend;
182        config.proxy_url = proxy_url.clone().map(SecretString::from);
183
184        Self {
185            url: url.into(),
186            config,
187            signal: Arc::new(AtomicBool::new(false)),
188            connection_mode,
189            cmd_tx: Arc::new(tokio::sync::RwLock::new(cmd_tx)),
190            out_rx: None,
191            handler_tasks: Arc::new(TaskGroup::new()),
192            connect_lock: Arc::new(tokio::sync::Mutex::new(())),
193            subscriptions: SubscriptionState::new(KRAKEN_SPOT_WS_TOPIC_DELIMITER),
194            subscription_payloads: Arc::new(tokio::sync::RwLock::new(HashMap::new())),
195            auth_tracker: AuthTracker::new(),
196            cancellation_token,
197            req_id_counter: Arc::new(AtomicU64::new(0)),
198            auth_token: Arc::new(tokio::sync::RwLock::new(None)),
199            account_id: Arc::new(RwLock::new(None)),
200            truncated_id_map: Arc::new(AtomicMap::new()),
201            instruments: Arc::new(AtomicMap::new()),
202            l2_depths: L2Depths::default(),
203            l2_book_commands: Arc::new(tokio::sync::Mutex::new(())),
204            book_requests: Arc::new(parking_lot::Mutex::new(ahash::AHashMap::new())),
205            l3_depths: Arc::new(parking_lot::Mutex::new(ahash::AHashMap::new())),
206            transport_backend,
207            proxy_url: proxy_url.map(SecretString::from),
208            socket_control: None,
209        }
210    }
211
212    pub(crate) fn begin_shutdown(&self) {
213        self.handler_tasks.begin_shutdown();
214        self.cancellation_token.cancel();
215        self.signal.store(true, Ordering::Relaxed);
216    }
217
218    /// Configures socket state reporting and reconnect control.
219    #[must_use]
220    pub fn with_socket_control(mut self, control: SocketControl) -> Self {
221        self.socket_control = Some(control);
222        self
223    }
224
225    fn get_next_req_id(&self) -> u64 {
226        self.req_id_counter.fetch_add(1, Ordering::Relaxed) + 1
227    }
228
229    /// Returns the shared request-id counter.
230    pub fn req_id_counter(&self) -> Arc<AtomicU64> {
231        self.req_id_counter.clone()
232    }
233
234    /// Returns a clone of the handler command channel sender.
235    pub async fn handler_command_sender(
236        &self,
237    ) -> tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand> {
238        self.cmd_tx.read().await.clone()
239    }
240
241    /// Returns the shared `cmd_tx` handle. Unlike
242    /// [`handler_command_sender`](Self::handler_command_sender) (a snapshot
243    /// clone), this exposes the `RwLock` so callers see the live sender
244    /// after `connect()` swaps it in.
245    pub fn handler_command_handle(
246        &self,
247    ) -> Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<SpotHandlerCommand>>> {
248        self.cmd_tx.clone()
249    }
250
251    /// Returns the current cached authentication token, if any.
252    pub async fn auth_token(&self) -> Option<SecretString> {
253        self.auth_token.read().await.clone()
254    }
255
256    /// Returns the current cached authentication token without awaiting.
257    ///
258    /// Returns `None` when the lock is contended or no token is cached. Used by
259    /// the synchronous order-routing path where the auth token is normally
260    /// uncontended; callers fall back to REST when the lock is unavailable.
261    pub fn auth_token_blocking(&self) -> Option<SecretString> {
262        self.auth_token
263            .try_read()
264            .ok()
265            .and_then(|guard| guard.clone())
266    }
267
268    /// Returns the shared auth token handle for internal components that need
269    /// non-async access, such as timeout-triggered compensating cancels.
270    pub(crate) fn auth_token_handle(&self) -> Arc<tokio::sync::RwLock<Option<SecretString>>> {
271        self.auth_token.clone()
272    }
273
274    /// Connects to the WebSocket server.
275    pub async fn connect(&mut self) -> Result<(), KrakenWsError> {
276        let connect_lock = Arc::clone(&self.connect_lock);
277        let _connect_guard = connect_lock.lock().await;
278
279        log::debug!("Connecting to {}", self.url.origin());
280
281        if !self.handler_tasks.is_open() || !self.handler_tasks.is_empty() {
282            self.disconnect_locked().await?;
283            self.handler_tasks.start_generation().map_err(|e| {
284                KrakenWsError::ConnectionError(format!(
285                    "Failed to start WebSocket handler task generation: {e}"
286                ))
287            })?;
288        }
289        let handler_spawner = self.handler_tasks.spawner().map_err(|e| {
290            KrakenWsError::ConnectionError(format!(
291                "Failed to acquire WebSocket handler task spawner: {e}"
292            ))
293        })?;
294
295        if self.cancellation_token.is_cancelled() {
296            self.cancellation_token = CancellationToken::new();
297        }
298
299        self.signal.store(false, Ordering::Relaxed);
300
301        let (raw_handler, raw_rx) = channel_message_handler();
302        let headers = create_standard_nautilus_headers();
303
304        let ws_config = WebSocketConfig {
305            url: self.url.expose().to_owned(),
306            headers,
307            heartbeat_interval_secs: Some(self.config.heartbeat_interval_secs),
308            heartbeat_payload: Some(WS_PING_MSG.to_string()),
309            connect_timeout_ms: Some(5_000),
310            reconnect_delay_initial_ms: Some(500),
311            reconnect_delay_max_ms: Some(5_000),
312            reconnect_backoff_factor: Some(1.5),
313            reconnect_jitter_ms: Some(250),
314            reconnect_max_attempts: None,
315            // Treat a silent connection as dead so the reconnect + resubscribe
316            // path runs. `0` disables; see `ws_idle_timeout_ms` docs (issue #4255).
317            heartbeat_timeout_secs: None,
318            idle_timeout_ms: (self.config.ws_idle_timeout_ms != 0)
319                .then_some(self.config.ws_idle_timeout_ms),
320            writer_capacity: None,
321            backend: self.transport_backend,
322            proxy_url: self
323                .proxy_url
324                .as_ref()
325                .map(|value| value.expose_secret().to_owned()),
326            max_message_size_bytes: None,
327            max_frame_size_bytes: None,
328        };
329
330        let keyed_quotas = vec![
331            (
332                KRAKEN_RATE_LIMIT_KEY_SUBSCRIPTION[0].to_string(),
333                *KRAKEN_SPOT_WS_SUBSCRIPTION_QUOTA,
334            ),
335            (
336                KRAKEN_RATE_LIMIT_KEY_ORDER[0].to_string(),
337                *KRAKEN_SPOT_WS_ORDER_QUOTA,
338            ),
339        ];
340
341        let ws_client = WebSocketClient::builder()
342            .config(ws_config)
343            .message_handler(raw_handler)
344            .keyed_quotas(keyed_quotas)
345            .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
346            .connect()
347            .await
348            .map_err(|e| KrakenWsError::ConnectionError(e.to_string()))?;
349
350        // Share connection state across clones via ArcSwap
351        self.connection_mode
352            .store(ws_client.connection_mode_atomic());
353        let reconnect_handle = ws_client.reconnect_handle();
354
355        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<KrakenSpotWsMessage>();
356        self.out_rx = Some(Arc::new(out_rx));
357
358        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<SpotHandlerCommand>();
359        *self.cmd_tx.write().await = cmd_tx.clone();
360
361        if let Err(e) = cmd_tx.send(SpotHandlerCommand::SetClient(ws_client)) {
362            return Err(KrakenWsError::ConnectionError(format!(
363                "Failed to send WebSocketClient to handler: {e}"
364            )));
365        }
366
367        if let Some(control) = &self.socket_control {
368            control.register(move || reconnect_handle.request_reconnect());
369        }
370
371        let signal = self.signal.clone();
372        let subscriptions = self.subscriptions.clone();
373        let subscription_payloads = self.subscription_payloads.clone();
374        let config_for_reconnect = self.config.clone();
375        let auth_token_for_reconnect = self.auth_token.clone();
376        let auth_tracker_for_reconnect = self.auth_tracker.clone();
377        let cmd_tx_for_reconnect = cmd_tx.clone();
378        let book_requests_for_reconnect = Arc::clone(&self.book_requests);
379        let l2_depths_for_reconnect = self.l2_depths.clone();
380
381        let handler_task = async move {
382            let mut handler =
383                SpotFeedHandler::new(signal.clone(), cmd_rx, raw_rx, subscriptions.clone());
384
385            loop {
386                match handler.next().await {
387                    Some(KrakenSpotWsMessage::Reconnected) => {
388                        if signal.load(Ordering::Relaxed) {
389                            continue;
390                        }
391                        log::info!("WebSocket reconnected, resubscribing");
392
393                        begin_resubscribe(
394                            &subscriptions,
395                            &book_requests_for_reconnect,
396                            &l2_depths_for_reconnect,
397                        );
398
399                        let payloads = subscription_payloads.read().await;
400                        if payloads.is_empty() {
401                            log::debug!("No subscriptions to restore after reconnection");
402                        } else {
403                            let had_auth = auth_token_for_reconnect.read().await.is_some();
404
405                            if had_auth && config_for_reconnect.has_api_credentials() {
406                                log::debug!("Re-authenticating after reconnect");
407
408                                auth_tracker_for_reconnect.invalidate();
409                                let _rx = auth_tracker_for_reconnect.begin();
410
411                                match refresh_auth_token(&config_for_reconnect).await {
412                                    Ok(new_token) => {
413                                        *auth_token_for_reconnect.write().await = Some(new_token);
414                                        auth_tracker_for_reconnect.succeed();
415                                        log::debug!("Re-authentication successful");
416                                    }
417                                    Err(e) => {
418                                        log::error!(
419                                            "Failed to re-authenticate after reconnect: {e}"
420                                        );
421                                        *auth_token_for_reconnect.write().await = None;
422                                        auth_tracker_for_reconnect.fail(e.to_string());
423                                    }
424                                }
425                            }
426
427                            log::debug!(
428                                "Resubscribing after reconnection: count={}",
429                                payloads.len()
430                            );
431
432                            for (topic, payload) in payloads.iter() {
433                                let needs_token =
434                                    topic == "executions" || topic.starts_with("level3:");
435                                let payload = if needs_token {
436                                    let auth_token = auth_token_for_reconnect.read().await.clone();
437                                    match auth_token {
438                                        Some(token) => {
439                                            match update_auth_token_in_payload(
440                                                payload.expose_secret(),
441                                                token.expose_secret(),
442                                            ) {
443                                                Ok(p) => p,
444                                                Err(e) => {
445                                                    log::error!("Failed to update auth token: {e}");
446                                                    continue;
447                                                }
448                                            }
449                                        }
450                                        None => {
451                                            log::warn!(
452                                                "Cannot resubscribe to {topic}: no auth token"
453                                            );
454                                            continue;
455                                        }
456                                    }
457                                } else {
458                                    payload.clone()
459                                };
460
461                                if let Err(e) = cmd_tx_for_reconnect
462                                    .send(SpotHandlerCommand::Subscribe { payload })
463                                {
464                                    log::error!(
465                                        "Failed to send resubscribe command: error={e}, \
466                                        topic={topic}"
467                                    );
468                                }
469
470                                subscriptions.mark_subscribe(topic);
471                            }
472                        }
473
474                        if out_tx.send(KrakenSpotWsMessage::Reconnected).is_err() {
475                            if handler.is_stopped() {
476                                log::debug!("Failed to send message (receiver dropped)");
477                            } else {
478                                log::error!("Failed to send message (receiver dropped)");
479                            }
480                            break;
481                        }
482                    }
483                    Some(msg) => {
484                        if out_tx.send(msg).is_err() {
485                            if handler.is_stopped() {
486                                log::debug!("Failed to send message (receiver dropped)");
487                            } else {
488                                log::error!("Failed to send message (receiver dropped)");
489                            }
490                            break;
491                        }
492                    }
493                    None => {
494                        if handler.is_stopped() {
495                            log::debug!("Stop signal received, ending message processing");
496                            break;
497                        }
498                        log::warn!("WebSocket stream ended unexpectedly");
499                        break;
500                    }
501                }
502            }
503
504            log::debug!("Handler task exiting");
505        };
506
507        if let Err(e) = handler_spawner.spawn(handler_task) {
508            if let Some(control) = &self.socket_control {
509                control.deregister();
510            }
511            self.out_rx = None;
512            return Err(KrakenWsError::ConnectionError(format!(
513                "Failed to register WebSocket handler task: {e}"
514            )));
515        }
516
517        log::debug!("WebSocket connected successfully");
518        Ok(())
519    }
520
521    /// Disconnects from the WebSocket server.
522    pub async fn disconnect(&mut self) -> Result<(), KrakenWsError> {
523        let connect_lock = Arc::clone(&self.connect_lock);
524        let _connect_guard = connect_lock.lock().await;
525        self.disconnect_locked().await
526    }
527
528    async fn disconnect_locked(&self) -> Result<(), KrakenWsError> {
529        log::debug!("Disconnecting WebSocket");
530
531        self.handler_tasks.begin_shutdown();
532        self.signal.store(true, Ordering::Relaxed);
533
534        if let Err(e) = self
535            .cmd_tx
536            .read()
537            .await
538            .send(SpotHandlerCommand::Disconnect)
539        {
540            log::debug!(
541                "Failed to send disconnect command (handler may already be shut down): {e}"
542            );
543        }
544
545        let task_result = self
546            .handler_tasks
547            .finish_shutdown(
548                tokio::time::Duration::from_secs(2),
549                tokio::time::Duration::from_secs(2),
550            )
551            .await;
552
553        self.subscriptions.clear();
554        self.subscription_payloads.write().await.clear();
555        self.auth_tracker.fail("Disconnected");
556
557        self.l3_depths.lock().clear();
558        self.l2_depths.clear();
559        self.book_requests.lock().clear();
560
561        if let Some(control) = &self.socket_control {
562            control.deregister();
563        }
564
565        match task_result {
566            Ok(()) => Ok(()),
567            Err(error @ TaskShutdownError::Timeout { .. }) => Err(KrakenWsError::Timeout(format!(
568                "Spot WebSocket handler shutdown timed out: {error}"
569            ))),
570            Err(e) => Err(KrakenWsError::Disconnected(format!(
571                "Spot WebSocket handler shutdown failed: {e}"
572            ))),
573        }
574    }
575
576    /// Closes the WebSocket connection.
577    pub async fn close(&mut self) -> Result<(), KrakenWsError> {
578        self.disconnect().await
579    }
580
581    /// Waits until the connection is active or timeout.
582    pub async fn wait_until_active(&self, timeout_secs: f64) -> Result<(), KrakenWsError> {
583        let timeout = tokio::time::Duration::from_secs_f64(timeout_secs);
584
585        tokio::time::timeout(timeout, async {
586            while !self.is_active() {
587                tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
588            }
589        })
590        .await
591        .map_err(|_| {
592            KrakenWsError::ConnectionError(format!(
593                "WebSocket connection timeout after {timeout_secs} seconds"
594            ))
595        })?;
596
597        Ok(())
598    }
599
600    /// Returns true if the WebSocket is authenticated for private subscriptions.
601    #[must_use]
602    pub fn is_authenticated(&self) -> bool {
603        self.auth_tracker.is_authenticated()
604    }
605
606    /// Waits until the WebSocket is authenticated or the timeout elapses.
607    ///
608    /// Returns an error on timeout or explicit auth failure.
609    pub async fn wait_until_authenticated(&self, timeout_secs: f64) -> Result<(), KrakenWsError> {
610        let timeout = tokio::time::Duration::from_secs_f64(timeout_secs);
611
612        if self.auth_tracker.wait_for_authenticated(timeout).await {
613            Ok(())
614        } else {
615            Err(KrakenWsError::AuthenticationError(format!(
616                "Authentication not completed within {timeout_secs} seconds"
617            )))
618        }
619    }
620
621    /// Authenticates with the Kraken API to enable private subscriptions.
622    pub async fn authenticate(&self) -> Result<(), KrakenWsError> {
623        if !self.config.has_api_credentials() {
624            return Err(KrakenWsError::AuthenticationError(
625                "API credentials required for authentication".to_string(),
626            ));
627        }
628
629        let _receiver = self.auth_tracker.begin();
630
631        match refresh_auth_token(&self.config).await {
632            Ok(token) => {
633                *self.auth_token.write().await = Some(token);
634                self.auth_tracker.succeed();
635                Ok(())
636            }
637            Err(e) => {
638                *self.auth_token.write().await = None;
639                self.auth_tracker.fail(e.to_string());
640                Err(e)
641            }
642        }
643    }
644
645    /// Cancels all pending requests.
646    pub fn cancel_all_requests(&self) {
647        self.cancellation_token.cancel();
648    }
649
650    /// Returns the cancellation token for this client.
651    pub fn cancellation_token(&self) -> &CancellationToken {
652        &self.cancellation_token
653    }
654
655    /// Subscribes to a channel for the given symbols.
656    pub async fn subscribe(
657        &self,
658        channel: KrakenWsChannel,
659        symbols: Vec<Ustr>,
660        depth: Option<u32>,
661    ) -> Result<(), KrakenWsError> {
662        if matches!(channel, KrakenWsChannel::Level3) {
663            return Err(KrakenWsError::InvalidMessage(
664                "Use subscribe_book_l3 / unsubscribe_book_l3 for the Level3 channel".to_string(),
665            ));
666        }
667        let mut symbols_to_subscribe = Vec::new();
668        let channel_str = channel.as_ref();
669        for symbol in &symbols {
670            let key = format!("{channel_str}:{symbol}");
671            if self.subscriptions.add_reference(&key) {
672                self.subscriptions.mark_subscribe(&key);
673                symbols_to_subscribe.push(*symbol);
674            }
675        }
676
677        if symbols_to_subscribe.is_empty() {
678            return Ok(());
679        }
680
681        let is_private = matches!(
682            channel,
683            KrakenWsChannel::Executions | KrakenWsChannel::Balances
684        );
685        let token = if is_private {
686            Some(self.auth_token().await.ok_or_else(|| {
687                KrakenWsError::AuthenticationError(
688                    "Authentication token required for private channels. Call authenticate() first"
689                        .to_string(),
690                )
691            })?)
692        } else {
693            None
694        };
695
696        let req_id = self.get_next_req_id();
697        let request = KrakenWsRequest {
698            method: KrakenWsMethod::Subscribe,
699            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
700                channel,
701                symbol: Some(symbols_to_subscribe.clone()),
702                snapshot: None,
703                depth,
704                interval: None,
705                event_trigger: None,
706                token,
707                snap_orders: None,
708                snap_trades: None,
709            })),
710            req_id: Some(req_id),
711        };
712
713        let payload = self.send_command(&request).await?;
714
715        for symbol in &symbols_to_subscribe {
716            let key = format!("{channel_str}:{symbol}");
717            self.subscriptions.confirm_subscribe(&key);
718            self.subscription_payloads
719                .write()
720                .await
721                .insert(key, payload.clone());
722        }
723
724        Ok(())
725    }
726
727    /// Subscribes to a channel with a specific interval (for OHLC).
728    async fn subscribe_with_interval(
729        &self,
730        channel: KrakenWsChannel,
731        symbols: Vec<Ustr>,
732        interval: u32,
733    ) -> Result<(), KrakenWsError> {
734        let mut symbols_to_subscribe = Vec::new();
735        let channel_str = channel.as_ref();
736        for symbol in &symbols {
737            let key = format!("{channel_str}:{symbol}:{interval}");
738            if self.subscriptions.add_reference(&key) {
739                self.subscriptions.mark_subscribe(&key);
740                symbols_to_subscribe.push(*symbol);
741            }
742        }
743
744        if symbols_to_subscribe.is_empty() {
745            return Ok(());
746        }
747
748        let req_id = self.get_next_req_id();
749        let request = KrakenWsRequest {
750            method: KrakenWsMethod::Subscribe,
751            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
752                channel,
753                symbol: Some(symbols_to_subscribe.clone()),
754                snapshot: Some(false),
755                depth: None,
756                interval: Some(interval),
757                event_trigger: None,
758                token: None,
759                snap_orders: None,
760                snap_trades: None,
761            })),
762            req_id: Some(req_id),
763        };
764
765        let payload = self.send_command(&request).await?;
766
767        for symbol in &symbols_to_subscribe {
768            let key = format!("{channel_str}:{symbol}:{interval}");
769            self.subscriptions.confirm_subscribe(&key);
770            self.subscription_payloads
771                .write()
772                .await
773                .insert(key, payload.clone());
774        }
775
776        Ok(())
777    }
778
779    /// Unsubscribes from a channel with a specific interval (for OHLC).
780    async fn unsubscribe_with_interval(
781        &self,
782        channel: KrakenWsChannel,
783        symbols: Vec<Ustr>,
784        interval: u32,
785    ) -> Result<(), KrakenWsError> {
786        let mut symbols_to_unsubscribe = Vec::new();
787        let channel_str = channel.as_ref();
788        for symbol in &symbols {
789            let key = format!("{channel_str}:{symbol}:{interval}");
790            if self.subscriptions.remove_reference(&key) {
791                self.subscriptions.mark_unsubscribe(&key);
792                symbols_to_unsubscribe.push(*symbol);
793            }
794        }
795
796        if symbols_to_unsubscribe.is_empty() {
797            return Ok(());
798        }
799
800        let req_id = self.get_next_req_id();
801        let request = KrakenWsRequest {
802            method: KrakenWsMethod::Unsubscribe,
803            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
804                channel,
805                symbol: Some(symbols_to_unsubscribe.clone()),
806                snapshot: None,
807                depth: None,
808                interval: Some(interval),
809                event_trigger: None,
810                token: None,
811                snap_orders: None,
812                snap_trades: None,
813            })),
814            req_id: Some(req_id),
815        };
816
817        self.send_command(&request).await?;
818
819        for symbol in &symbols_to_unsubscribe {
820            let key = format!("{channel_str}:{symbol}:{interval}");
821            self.subscriptions.confirm_unsubscribe(&key);
822            self.subscription_payloads.write().await.remove(&key);
823        }
824
825        Ok(())
826    }
827
828    /// Unsubscribes from a channel for the given symbols.
829    pub async fn unsubscribe(
830        &self,
831        channel: KrakenWsChannel,
832        symbols: Vec<Ustr>,
833    ) -> Result<(), KrakenWsError> {
834        if matches!(channel, KrakenWsChannel::Level3) {
835            return Err(KrakenWsError::InvalidMessage(
836                "Use subscribe_book_l3 / unsubscribe_book_l3 for the Level3 channel".to_string(),
837            ));
838        }
839        let mut symbols_to_unsubscribe = Vec::new();
840        let channel_str = channel.as_ref();
841        for symbol in &symbols {
842            let key = format!("{channel_str}:{symbol}");
843            if self.subscriptions.remove_reference(&key) {
844                self.subscriptions.mark_unsubscribe(&key);
845                symbols_to_unsubscribe.push(*symbol);
846            } else {
847                log::debug!(
848                    "Channel {channel_str} symbol {symbol} still has active subscriptions, not unsubscribing"
849                );
850            }
851        }
852
853        if symbols_to_unsubscribe.is_empty() {
854            return Ok(());
855        }
856
857        let is_private = matches!(
858            channel,
859            KrakenWsChannel::Executions | KrakenWsChannel::Balances
860        );
861        let token = if is_private {
862            Some(self.auth_token().await.ok_or_else(|| {
863                KrakenWsError::AuthenticationError(
864                    "Authentication token required for private channels. Call authenticate() first"
865                        .to_string(),
866                )
867            })?)
868        } else {
869            None
870        };
871
872        let req_id = self.get_next_req_id();
873        let request = KrakenWsRequest {
874            method: KrakenWsMethod::Unsubscribe,
875            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
876                channel,
877                symbol: Some(symbols_to_unsubscribe.clone()),
878                snapshot: None,
879                depth: None,
880                interval: None,
881                event_trigger: None,
882                token,
883                snap_orders: None,
884                snap_trades: None,
885            })),
886            req_id: Some(req_id),
887        };
888
889        self.send_command(&request).await?;
890
891        for symbol in &symbols_to_unsubscribe {
892            let key = format!("{channel_str}:{symbol}");
893            self.subscriptions.confirm_unsubscribe(&key);
894            self.subscription_payloads.write().await.remove(&key);
895        }
896
897        Ok(())
898    }
899
900    /// Sends a ping message to keep the connection alive.
901    pub async fn send_ping(&self) -> Result<(), KrakenWsError> {
902        let req_id = self.get_next_req_id();
903
904        let request = KrakenWsRequest {
905            method: KrakenWsMethod::Ping,
906            params: None,
907            req_id: Some(req_id),
908        };
909
910        self.send_command(&request).await?;
911        Ok(())
912    }
913
914    async fn send_command(&self, request: &KrakenWsRequest) -> Result<SecretString, KrakenWsError> {
915        let payload = SecretString::from(
916            serde_json::to_string(request).map_err(|e| KrakenWsError::JsonError(e.to_string()))?,
917        );
918
919        log::trace!(
920            "Sending WebSocket request: method={:?} ({} bytes)",
921            request.method,
922            payload.expose_secret().len(),
923        );
924
925        let cmd = match request.method {
926            KrakenWsMethod::Subscribe => SpotHandlerCommand::Subscribe {
927                payload: payload.clone(),
928            },
929            KrakenWsMethod::Unsubscribe => SpotHandlerCommand::Unsubscribe {
930                payload: payload.clone(),
931            },
932            KrakenWsMethod::Ping | KrakenWsMethod::Pong => SpotHandlerCommand::Ping {
933                payload: payload.clone(),
934            },
935            KrakenWsMethod::AddOrder
936            | KrakenWsMethod::AmendOrder
937            | KrakenWsMethod::CancelOrder
938            | KrakenWsMethod::BatchAdd => {
939                return Err(KrakenWsError::InvalidMessage(
940                    "Order methods must not be sent via send_command; use the dedicated order submission path".to_string()
941                ));
942            }
943        };
944
945        self.cmd_tx
946            .read()
947            .await
948            .send(cmd)
949            .map_err(|e| KrakenWsError::ConnectionError(format!("Failed to send request: {e}")))?;
950
951        Ok(payload)
952    }
953
954    /// Returns true if connected (not closed).
955    pub fn is_connected(&self) -> bool {
956        let connection_mode_arc = self.connection_mode.load();
957        !ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
958    }
959
960    /// Returns true if the connection is active.
961    pub fn is_active(&self) -> bool {
962        let connection_mode_arc = self.connection_mode.load();
963        ConnectionMode::from_atomic(&connection_mode_arc).is_active()
964            && !self.signal.load(Ordering::Relaxed)
965    }
966
967    /// Returns true if the connection is closed.
968    pub fn is_closed(&self) -> bool {
969        let connection_mode_arc = self.connection_mode.load();
970        ConnectionMode::from_atomic(&connection_mode_arc).is_closed()
971            || self.signal.load(Ordering::Relaxed)
972    }
973
974    /// Returns the WebSocket URL.
975    pub fn url(&self) -> &str {
976        self.url.expose()
977    }
978
979    /// Returns all active subscriptions.
980    pub fn get_subscriptions(&self) -> Vec<String> {
981        self.subscriptions.all_topics()
982    }
983
984    /// Returns `true` if a topic is currently subscribed (confirmed).
985    pub fn subscriptions_contains(&self, topic: &str) -> bool {
986        self.subscriptions.all_topics().iter().any(|t| t == topic)
987    }
988
989    /// Sets the account ID for execution report parsing.
990    pub fn set_account_id(&self, account_id: AccountId) {
991        *self.account_id.write() = Some(account_id);
992    }
993
994    /// Returns the account ID if set.
995    #[must_use]
996    pub fn account_id(&self) -> Option<AccountId> {
997        *self.account_id.read()
998    }
999
1000    /// Caches an instrument for execution report parsing.
1001    pub fn cache_instrument(&self, instrument: InstrumentAny) {
1002        self.instruments.insert(instrument.id(), instrument);
1003    }
1004
1005    /// Returns a shared reference to the account ID.
1006    pub fn account_id_shared(&self) -> &Arc<RwLock<Option<AccountId>>> {
1007        &self.account_id
1008    }
1009
1010    /// Returns a shared reference to the truncated ID map.
1011    pub fn truncated_id_map(&self) -> &Arc<AtomicMap<String, ClientOrderId>> {
1012        &self.truncated_id_map
1013    }
1014
1015    /// Caches a client order for truncated ID resolution.
1016    pub fn cache_client_order(
1017        &self,
1018        client_order_id: ClientOrderId,
1019        _venue_order_id: Option<VenueOrderId>,
1020        _instrument_id: InstrumentId,
1021        _trader_id: TraderId,
1022        _strategy_id: StrategyId,
1023    ) {
1024        let truncated = crate::common::parse::truncate_cl_ord_id(&client_order_id);
1025
1026        if truncated != client_order_id.as_str() {
1027            self.truncated_id_map.insert(truncated, client_order_id);
1028        }
1029    }
1030
1031    /// Returns a stream of WebSocket messages.
1032    ///
1033    /// # Errors
1034    ///
1035    /// Returns an error if:
1036    /// - The stream receiver has already been taken.
1037    /// - Other clones of this client still hold references to the receiver.
1038    pub fn stream(
1039        &mut self,
1040    ) -> Result<impl futures_util::Stream<Item = KrakenSpotWsMessage> + use<>, KrakenWsError> {
1041        let rx = self.out_rx.take().ok_or_else(|| {
1042            KrakenWsError::ChannelError(
1043                "Stream receiver already taken or client not connected".to_string(),
1044            )
1045        })?;
1046        let mut rx = Arc::try_unwrap(rx).map_err(|_| {
1047            KrakenWsError::ChannelError(
1048                "Cannot take ownership of stream - other client clones still hold references"
1049                    .to_string(),
1050            )
1051        })?;
1052        Ok(async_stream::stream! {
1053            while let Some(msg) = rx.recv().await {
1054                yield msg;
1055            }
1056        })
1057    }
1058
1059    /// Subscribes to order book updates for the given instrument.
1060    pub async fn subscribe_book(
1061        &self,
1062        instrument_id: InstrumentId,
1063        depth: Option<u32>,
1064    ) -> Result<(), KrakenWsError> {
1065        let _book_commands = self.l2_book_commands.lock().await;
1066        let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1067        let depth = depth.unwrap_or(10);
1068
1069        if !matches!(depth, 10 | 25 | 100 | 500 | 1000) {
1070            return Err(KrakenWsError::InvalidMessage(format!(
1071                "Invalid L2 depth {depth}, valid values: 10, 25, 100, 500, 1000",
1072            )));
1073        }
1074
1075        let channel_str = KrakenWsChannel::Book.as_ref();
1076        let key = format!("{channel_str}:{symbol}");
1077        let is_first_reference = self.subscriptions.add_reference(&key);
1078
1079        if !is_first_reference {
1080            let existing_depth = self.l2_depths.get(symbol.as_str());
1081
1082            if existing_depth != Some(depth) {
1083                self.subscriptions.remove_reference(&key);
1084                return Err(KrakenWsError::InvalidMessage(format!(
1085                    "L2 subscription for {symbol} already exists with depth \
1086                     {existing_depth:?}, cannot resubscribe with depth {depth}",
1087                )));
1088            }
1089            return Ok(());
1090        }
1091
1092        self.subscriptions.mark_subscribe(&key);
1093        let req_id = self.get_next_req_id();
1094        self.l2_depths.insert(symbol.as_str(), depth, req_id);
1095
1096        let request = KrakenWsRequest {
1097            method: KrakenWsMethod::Subscribe,
1098            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1099                channel: KrakenWsChannel::Book,
1100                symbol: Some(vec![symbol]),
1101                snapshot: None,
1102                depth: Some(depth),
1103                interval: None,
1104                event_trigger: None,
1105                token: None,
1106                snap_orders: None,
1107                snap_trades: None,
1108            })),
1109            req_id: Some(req_id),
1110        };
1111        self.record_book_request(req_id, symbol);
1112
1113        let payload = match self.send_command(&request).await {
1114            Ok(payload) => payload,
1115            Err(e) => {
1116                self.book_requests.lock().remove(&req_id);
1117                self.l2_depths.remove(symbol.as_str());
1118                self.subscriptions.remove_reference(&key);
1119                self.subscriptions.mark_unsubscribe(&key);
1120                self.subscriptions.confirm_unsubscribe(&key);
1121                return Err(e);
1122            }
1123        };
1124
1125        self.subscriptions.confirm_subscribe(&key);
1126        self.subscription_payloads
1127            .write()
1128            .await
1129            .insert(key, payload);
1130        Ok(())
1131    }
1132
1133    /// Subscribes to the `level3` channel for the given Kraken symbol.
1134    ///
1135    /// `depth` must be one of `10`, `100`, or `1000`. The depth is recorded in
1136    /// the per-client `l3_depths` map so the message handler and `resync_book_l3`
1137    /// can recover it without a separate side-table.
1138    ///
1139    /// If the symbol is already subscribed, the existing depth must match - Kraken
1140    /// streams one depth per `(symbol, channel)` pair, so a second subscribe with
1141    /// a different depth would corrupt the local runtime state. Mismatch returns
1142    /// an error without mutating state.
1143    ///
1144    /// # Errors
1145    ///
1146    /// Returns an error if `depth` is invalid, the auth token is not cached
1147    /// (call `authenticate` first), the requested depth differs from an existing
1148    /// subscription, or the message cannot be sent. On any of those failures
1149    /// the reference count and pending state are rolled back fully so callers
1150    /// retry from a clean state.
1151    pub async fn subscribe_book_l3(&self, symbol: Ustr, depth: u32) -> Result<(), KrakenWsError> {
1152        if !matches!(depth, 10 | 100 | 1000) {
1153            return Err(KrakenWsError::InvalidMessage(format!(
1154                "Invalid L3 depth {depth}, valid values: 10, 100, 1000",
1155            )));
1156        }
1157
1158        let token = self.auth_token().await.ok_or_else(|| {
1159            KrakenWsError::AuthenticationError(
1160                "Authentication token required for level3. Call authenticate() first".to_string(),
1161            )
1162        })?;
1163
1164        let channel_str = KrakenWsChannel::Level3.as_ref();
1165        let key = format!("{channel_str}:{symbol}");
1166
1167        let is_first_reference = self.subscriptions.add_reference(&key);
1168
1169        if !is_first_reference {
1170            let existing_depth = self.l3_depths.lock().get(symbol.as_str()).copied();
1171
1172            if existing_depth != Some(depth) {
1173                self.subscriptions.remove_reference(&key);
1174                return Err(KrakenWsError::InvalidMessage(format!(
1175                    "L3 subscription for {symbol} already exists with depth \
1176                     {existing_depth:?}, cannot resubscribe with depth {depth}",
1177                )));
1178            }
1179            return Ok(());
1180        }
1181
1182        self.subscriptions.mark_subscribe(&key);
1183
1184        self.l3_depths.lock().insert(symbol.to_string(), depth);
1185
1186        let req_id = self.get_next_req_id();
1187        let request = KrakenWsRequest {
1188            method: KrakenWsMethod::Subscribe,
1189            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1190                channel: KrakenWsChannel::Level3,
1191                symbol: Some(vec![symbol]),
1192                snapshot: Some(true),
1193                depth: Some(depth),
1194                interval: None,
1195                event_trigger: None,
1196                token: Some(token),
1197                snap_orders: None,
1198                snap_trades: None,
1199            })),
1200            req_id: Some(req_id),
1201        };
1202
1203        let payload = match self.send_command(&request).await {
1204            Ok(p) => p,
1205            Err(e) => {
1206                self.l3_depths.lock().remove(symbol.as_str());
1207                self.subscriptions.remove_reference(&key);
1208                self.subscriptions.mark_unsubscribe(&key);
1209                self.subscriptions.confirm_unsubscribe(&key);
1210                return Err(e);
1211            }
1212        };
1213
1214        self.subscriptions.confirm_subscribe(&key);
1215        self.subscription_payloads
1216            .write()
1217            .await
1218            .insert(key, payload);
1219        Ok(())
1220    }
1221
1222    /// Unsubscribes from the `level3` channel for the given Kraken symbol.
1223    ///
1224    /// Mirrors the existing `KrakenSpotWebSocketClient::unsubscribe` pattern:
1225    /// `SubscriptionState` is mutated optimistically before `send_command`. If
1226    /// the send fails, the local state reflects an unsubscribe that the venue
1227    /// never received; the next reconnect's payload replay will not include the
1228    /// topic (it was removed from `subscription_payloads`). Tightening this to
1229    /// roll-back-on-send-fail is a codebase-wide pattern change (every channel
1230    /// has it) and is out of scope for this PR - the L3 path matches the
1231    /// existing surface area rather than introducing an inconsistent improvement.
1232    ///
1233    /// # Errors
1234    ///
1235    /// Returns an error if the message cannot be sent.
1236    pub async fn unsubscribe_book_l3(&self, symbol: Ustr) -> Result<(), KrakenWsError> {
1237        let channel_str = KrakenWsChannel::Level3.as_ref();
1238        let key = format!("{channel_str}:{symbol}");
1239        if !self.subscriptions.remove_reference(&key) {
1240            return Ok(());
1241        }
1242        self.subscriptions.mark_unsubscribe(&key);
1243
1244        let token = self.auth_token().await;
1245        let req_id = self.get_next_req_id();
1246        let request = KrakenWsRequest {
1247            method: KrakenWsMethod::Unsubscribe,
1248            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1249                channel: KrakenWsChannel::Level3,
1250                symbol: Some(vec![symbol]),
1251                snapshot: None,
1252                depth: None,
1253                interval: None,
1254                event_trigger: None,
1255                token,
1256                snap_orders: None,
1257                snap_trades: None,
1258            })),
1259            req_id: Some(req_id),
1260        };
1261
1262        self.send_command(&request).await?;
1263        self.subscriptions.confirm_unsubscribe(&key);
1264        self.subscription_payloads.write().await.remove(&key);
1265        self.l3_depths.lock().remove(symbol.as_str());
1266        Ok(())
1267    }
1268
1269    /// Resynchronizes the `level3` book for `symbol` after a checksum mismatch.
1270    ///
1271    /// Refreshes the auth token unconditionally and issues a venue-level
1272    /// unsubscribe followed by a subscribe with `snapshot=true`, **bypassing
1273    /// `SubscriptionState` reference counts** so the user's logical
1274    /// subscription survives. The reference count is not changed; if multiple
1275    /// callers hold references to the same symbol, all of them continue to see
1276    /// the symbol as subscribed throughout the resync.
1277    ///
1278    /// # Errors
1279    ///
1280    /// Returns an error if the auth token cannot be refreshed or the
1281    /// unsubscribe/subscribe messages cannot be sent.
1282    pub async fn resync_book_l3(&self, symbol: Ustr, depth: u32) -> Result<(), KrakenWsError> {
1283        let channel_str = KrakenWsChannel::Level3.as_ref();
1284        let key = format!("{channel_str}:{symbol}");
1285
1286        // A resync can be retrying / awaiting a token refresh while the user
1287        // unsubscribes. Bail before mutating any state if the user no longer
1288        // holds a logical subscription, otherwise the venue-level subscribe
1289        // below would resurrect an orphaned stream that `SubscriptionState`
1290        // can never tear down.
1291        if !self.subscriptions_contains(&key) {
1292            log::debug!("Skipping L3 resync: subscription cancelled mid-retry, symbol={symbol}",);
1293            return Ok(());
1294        }
1295
1296        let new_token = refresh_auth_token(&self.config).await?;
1297        *self.auth_token.write().await = Some(new_token.clone());
1298
1299        // Re-check after the await - the user may have unsubscribed while we
1300        // were minting a fresh token.
1301        if !self.subscriptions_contains(&key) {
1302            log::debug!(
1303                "Skipping L3 resync: subscription cancelled after token refresh, symbol={symbol}",
1304            );
1305            return Ok(());
1306        }
1307
1308        let unsub_req_id = self.get_next_req_id();
1309        let unsub = KrakenWsRequest {
1310            method: KrakenWsMethod::Unsubscribe,
1311            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1312                channel: KrakenWsChannel::Level3,
1313                symbol: Some(vec![symbol]),
1314                snapshot: None,
1315                depth: None,
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(unsub_req_id),
1323        };
1324        self.send_command(&unsub).await?;
1325
1326        // Final check before issuing the resubscribe - same race window.
1327        if !self.subscriptions_contains(&key) {
1328            log::debug!("Skipping L3 resync resubscribe: cancelled before send, symbol={symbol}",);
1329            return Ok(());
1330        }
1331
1332        let sub_req_id = self.get_next_req_id();
1333        let sub = KrakenWsRequest {
1334            method: KrakenWsMethod::Subscribe,
1335            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1336                channel: KrakenWsChannel::Level3,
1337                symbol: Some(vec![symbol]),
1338                snapshot: Some(true),
1339                depth: Some(depth),
1340                interval: None,
1341                event_trigger: None,
1342                token: Some(new_token.clone()),
1343                snap_orders: None,
1344                snap_trades: None,
1345            })),
1346            req_id: Some(sub_req_id),
1347        };
1348        let payload = self.send_command(&sub).await?;
1349
1350        // Only persist replay payload + depth if the subscription is still
1351        // referenced. A late-arriving unsubscribe between `send_command` and
1352        // here would otherwise leave an orphan entry the reconnect path would
1353        // replay.
1354        if self.subscriptions_contains(&key) {
1355            self.subscription_payloads
1356                .write()
1357                .await
1358                .insert(key, payload);
1359            self.l3_depths.lock().insert(symbol.to_string(), depth);
1360        }
1361
1362        Ok(())
1363    }
1364
1365    /// Returns whether L2 `book` checksum validation is enabled for this client.
1366    pub fn validate_l2_checksum(&self) -> bool {
1367        self.config.validate_l2_checksum
1368    }
1369
1370    /// Resubscribes the `book` channel for `instrument_id` to recover its stream.
1371    ///
1372    /// The logical subscription is kept throughout, so the reconnect path still replays it.
1373    /// `generation` names the subscription the request belongs to and `epoch` the symbol's
1374    /// snapshot epoch when the request was issued. The check that this subscription is still the
1375    /// live one and the two commands form one step under the book command lock, which
1376    /// `subscribe_book` and `unsubscribe_book` also take, so a replacement or a cancel cannot land
1377    /// between them: a recovery for a replaced or canceled subscription sends nothing, and one that
1378    /// is admitted resubscribes at the live depth with an explicit snapshot, which the recovery
1379    /// depends on to end the wait.
1380    ///
1381    /// The epoch check and the recording of the resubscribe as the symbol's latest request are one
1382    /// compare-and-swap against the data client accepting a snapshot, so a late snapshot either
1383    /// serves the request, which then sends nothing, or is dropped as a retired stream's. The
1384    /// latest request stays recorded when a send fails: the data client then sees it unconfirmed
1385    /// and its snapshot watchdog asks again.
1386    ///
1387    /// # Errors
1388    ///
1389    /// Returns an error if either command fails to send.
1390    pub async fn resync_book(
1391        &self,
1392        instrument_id: InstrumentId,
1393        generation: u64,
1394        epoch: u64,
1395    ) -> Result<(), KrakenWsError> {
1396        let _book_commands = self.l2_book_commands.lock().await;
1397        let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1398        let channel_str = KrakenWsChannel::Book.as_ref();
1399        let key = format!("{channel_str}:{symbol}");
1400        let req_id = self.get_next_req_id();
1401        let live = self
1402            .subscriptions_contains(&key)
1403            .then(|| {
1404                self.l2_depths
1405                    .begin_resync(symbol.as_str(), generation, epoch, req_id)
1406            })
1407            .flatten();
1408
1409        let Some(live) = live else {
1410            // Nothing is withheld here, so the skip is not a warning: a canceled subscription has
1411            // no holder to serve, a replacement has a snapshot of its own on the way, a snapshot
1412            // accepted since the request has served it, and a reconnect replays the stream.
1413            log::debug!(
1414                "Skipping L2 resync: subscription canceled or replaced, or its snapshot served \
1415                 since, symbol={symbol}, generation={generation}, epoch={epoch}"
1416            );
1417            return Ok(());
1418        };
1419        let depth = live.depth;
1420
1421        let unsub = KrakenWsRequest {
1422            method: KrakenWsMethod::Unsubscribe,
1423            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1424                channel: KrakenWsChannel::Book,
1425                symbol: Some(vec![symbol]),
1426                snapshot: None,
1427                // The venue keys a `book` subscription by symbol and depth, and takes an
1428                // unsubscribe without a depth as depth 10.
1429                depth: Some(depth),
1430                interval: None,
1431                event_trigger: None,
1432                token: None,
1433                snap_orders: None,
1434                snap_trades: None,
1435            })),
1436            req_id: Some(self.get_next_req_id()),
1437        };
1438        self.send_command(&unsub).await?;
1439
1440        let sub = KrakenWsRequest {
1441            method: KrakenWsMethod::Subscribe,
1442            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1443                channel: KrakenWsChannel::Book,
1444                symbol: Some(vec![symbol]),
1445                // Asked for explicitly: the recovery depends on this snapshot to end the wait.
1446                snapshot: Some(true),
1447                depth: Some(depth),
1448                interval: None,
1449                event_trigger: None,
1450                token: None,
1451                snap_orders: None,
1452                snap_trades: None,
1453            })),
1454            req_id: Some(req_id),
1455        };
1456        self.record_book_request(req_id, symbol);
1457
1458        let payload = match self.send_command(&sub).await {
1459            Ok(payload) => payload,
1460            Err(e) => {
1461                self.book_requests.lock().remove(&req_id);
1462                return Err(e);
1463            }
1464        };
1465        self.subscription_payloads
1466            .write()
1467            .await
1468            .insert(key, payload);
1469
1470        Ok(())
1471    }
1472
1473    fn record_book_request(&self, req_id: u64, symbol: Ustr) {
1474        self.book_requests
1475            .lock()
1476            .insert(req_id, L2BookRequest { symbol });
1477    }
1478
1479    /// Returns whether L3 checksum validation is enabled for this client.
1480    pub fn validate_l3_checksum(&self) -> bool {
1481        self.config.validate_l3_checksum
1482    }
1483
1484    /// Returns `true` if the client has API credentials configured
1485    /// (post-environment-variable resolution).
1486    pub fn has_credentials(&self) -> bool {
1487        self.config.has_api_credentials()
1488    }
1489
1490    /// Returns a shared handle to the per-client instrument map.
1491    ///
1492    /// L3 stream-loop consumers read instruments through this handle so that
1493    /// `cache_instrument()` updates made after `connect()` are observed by the
1494    /// runtime book reconstruction without needing a re-connect.
1495    pub fn instruments_handle(&self) -> Arc<AtomicMap<InstrumentId, InstrumentAny>> {
1496        Arc::clone(&self.instruments)
1497    }
1498
1499    /// Returns a shared handle to the per-symbol L3 depth map.
1500    ///
1501    /// Stream-loop consumers read this map to drive `process_l3_message`'s
1502    /// resync depth lookup; `subscribe_book_l3` writes the depth.
1503    pub fn l3_depths_handle(&self) -> Arc<parking_lot::Mutex<ahash::AHashMap<String, u32>>> {
1504        Arc::clone(&self.l3_depths)
1505    }
1506
1507    /// Returns a shared handle to the per-symbol L2 depth map.
1508    pub(crate) fn l2_depths_handle(&self) -> L2Depths {
1509        self.l2_depths.clone()
1510    }
1511
1512    /// Returns a shared handle to the `book` subscribes awaiting the venue's answer.
1513    pub(crate) fn book_requests_handle(&self) -> L2BookRequests {
1514        Arc::clone(&self.book_requests)
1515    }
1516
1517    /// Subscribes to quote updates for the given instrument.
1518    ///
1519    /// Uses the Ticker channel with `event_trigger: "bbo"` for updates only on
1520    /// best bid/offer changes.
1521    pub async fn subscribe_quotes(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1522        let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1523        let key = format!("quotes:{symbol}");
1524
1525        if !self.subscriptions.add_reference(&key) {
1526            return Ok(());
1527        }
1528
1529        self.subscriptions.mark_subscribe(&key);
1530
1531        let req_id = self.get_next_req_id();
1532        let request = KrakenWsRequest {
1533            method: KrakenWsMethod::Subscribe,
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        let payload = self.send_command(&request).await?;
1549        self.subscriptions.confirm_subscribe(&key);
1550        self.subscription_payloads
1551            .write()
1552            .await
1553            .insert(key, payload);
1554        Ok(())
1555    }
1556
1557    /// Subscribes to trade updates for the given instrument.
1558    pub async fn subscribe_trades(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1559        let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1560        self.subscribe(KrakenWsChannel::Trade, vec![symbol], None)
1561            .await
1562    }
1563
1564    /// Subscribes to bar/OHLC updates for the given bar type.
1565    ///
1566    /// # Errors
1567    ///
1568    /// Returns an error if the bar aggregation is not supported by Kraken.
1569    pub async fn subscribe_bars(&self, bar_type: BarType) -> Result<(), KrakenWsError> {
1570        let symbol = to_ws_v2_symbol(bar_type.instrument_id().symbol.inner());
1571        let interval = bar_type_to_ws_interval(bar_type)?;
1572        self.subscribe_with_interval(KrakenWsChannel::Ohlc, vec![symbol], interval)
1573            .await
1574    }
1575
1576    /// Subscribes to execution updates (order and fill events).
1577    ///
1578    /// Requires authentication - call `authenticate()` first.
1579    pub async fn subscribe_executions(
1580        &self,
1581        snap_orders: bool,
1582        snap_trades: bool,
1583    ) -> Result<(), KrakenWsError> {
1584        let req_id = self.get_next_req_id();
1585
1586        let token = self.auth_token().await.ok_or_else(|| {
1587            KrakenWsError::AuthenticationError(
1588                "Authentication token required for executions channel. Call authenticate() first"
1589                    .to_string(),
1590            )
1591        })?;
1592
1593        let request = KrakenWsRequest {
1594            method: KrakenWsMethod::Subscribe,
1595            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1596                channel: KrakenWsChannel::Executions,
1597                symbol: None,
1598                snapshot: None,
1599                depth: None,
1600                interval: None,
1601                event_trigger: None,
1602                token: Some(token),
1603                snap_orders: Some(snap_orders),
1604                snap_trades: Some(snap_trades),
1605            })),
1606            req_id: Some(req_id),
1607        };
1608
1609        let payload = self.send_command(&request).await?;
1610
1611        let key = "executions";
1612        if self.subscriptions.add_reference(key) {
1613            self.subscriptions.mark_subscribe(key);
1614            self.subscriptions.confirm_subscribe(key);
1615            self.subscription_payloads
1616                .write()
1617                .await
1618                .insert(key.to_string(), payload);
1619        }
1620
1621        Ok(())
1622    }
1623
1624    /// Unsubscribes from order book updates for the given instrument.
1625    ///
1626    /// The unsubscribe carries the recorded depth: the venue keys a `book` subscription by symbol
1627    /// and depth, and takes an unsubscribe without a depth as depth 10.
1628    pub async fn unsubscribe_book(&self, instrument_id: InstrumentId) -> Result<(), KrakenWsError> {
1629        let _book_commands = self.l2_book_commands.lock().await;
1630        let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1631        let channel_str = KrakenWsChannel::Book.as_ref();
1632        let key = format!("{channel_str}:{symbol}");
1633
1634        if !self.subscriptions.remove_reference(&key) {
1635            return Ok(());
1636        }
1637
1638        self.subscriptions.mark_unsubscribe(&key);
1639
1640        let req_id = self.get_next_req_id();
1641        let request = KrakenWsRequest {
1642            method: KrakenWsMethod::Unsubscribe,
1643            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1644                channel: KrakenWsChannel::Book,
1645                symbol: Some(vec![symbol]),
1646                snapshot: None,
1647                depth: self.l2_depths.get(symbol.as_str()),
1648                interval: None,
1649                event_trigger: None,
1650                token: None,
1651                snap_orders: None,
1652                snap_trades: None,
1653            })),
1654            req_id: Some(req_id),
1655        };
1656
1657        self.send_command(&request).await?;
1658        self.subscriptions.confirm_unsubscribe(&key);
1659        self.subscription_payloads.write().await.remove(&key);
1660        self.l2_depths.remove(symbol.as_str());
1661        Ok(())
1662    }
1663
1664    /// Unsubscribes from quote updates for the given instrument.
1665    pub async fn unsubscribe_quotes(
1666        &self,
1667        instrument_id: InstrumentId,
1668    ) -> Result<(), KrakenWsError> {
1669        let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1670        let key = format!("quotes:{symbol}");
1671
1672        if !self.subscriptions.remove_reference(&key) {
1673            return Ok(());
1674        }
1675
1676        self.subscriptions.mark_unsubscribe(&key);
1677
1678        let req_id = self.get_next_req_id();
1679        let request = KrakenWsRequest {
1680            method: KrakenWsMethod::Unsubscribe,
1681            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1682                channel: KrakenWsChannel::Ticker,
1683                symbol: Some(vec![symbol]),
1684                snapshot: None,
1685                depth: None,
1686                interval: None,
1687                event_trigger: Some("bbo".to_string()),
1688                token: None,
1689                snap_orders: None,
1690                snap_trades: None,
1691            })),
1692            req_id: Some(req_id),
1693        };
1694
1695        self.send_command(&request).await?;
1696        self.subscriptions.confirm_unsubscribe(&key);
1697        self.subscription_payloads.write().await.remove(&key);
1698        Ok(())
1699    }
1700
1701    /// Unsubscribes from trade updates for the given instrument.
1702    pub async fn unsubscribe_trades(
1703        &self,
1704        instrument_id: InstrumentId,
1705    ) -> Result<(), KrakenWsError> {
1706        let symbol = to_ws_v2_symbol(instrument_id.symbol.inner());
1707        self.unsubscribe(KrakenWsChannel::Trade, vec![symbol]).await
1708    }
1709
1710    /// Unsubscribes from bar/OHLC updates for the given bar type.
1711    ///
1712    /// # Errors
1713    ///
1714    /// Returns an error if the bar aggregation is not supported by Kraken.
1715    pub async fn unsubscribe_bars(&self, bar_type: BarType) -> Result<(), KrakenWsError> {
1716        let symbol = to_ws_v2_symbol(bar_type.instrument_id().symbol.inner());
1717        let interval = bar_type_to_ws_interval(bar_type)?;
1718        self.unsubscribe_with_interval(KrakenWsChannel::Ohlc, vec![symbol], interval)
1719            .await
1720    }
1721}
1722
1723/// Opens a reconnect replay: the subscription state forgets its confirmations, the `book`
1724/// subscribes on record are retired and every `book` subscription's snapshot epoch moves.
1725///
1726/// The replay re-sends each stored payload under its original request id, so an answer to a replay
1727/// must not be read as the answer to the request that first sent it, and a request the dropped
1728/// socket never answered must not stay on record for the connection's lifetime. A replay the venue
1729/// rejects is left to the data client's snapshot watchdog. A recovery issued before the reconnect
1730/// finds the epoch moved and sends nothing, rather than cycling the replayed stream.
1731fn begin_resubscribe(
1732    subscriptions: &SubscriptionState,
1733    book_requests: &L2BookRequests,
1734    l2_depths: &L2Depths,
1735) {
1736    subscriptions.reset_after_reconnect();
1737    book_requests.lock().clear();
1738    l2_depths.advance_epochs();
1739}
1740
1741/// Refreshes the authentication token via the HTTP API.
1742async fn refresh_auth_token(
1743    config: &KrakenDataClientConfig,
1744) -> Result<SecretString, KrakenWsError> {
1745    let api_key = config
1746        .api_key
1747        .clone()
1748        .ok_or_else(|| KrakenWsError::AuthenticationError("Missing API key".to_string()))?;
1749    let api_secret = config
1750        .api_secret
1751        .clone()
1752        .ok_or_else(|| KrakenWsError::AuthenticationError("Missing API secret".to_string()))?;
1753
1754    let http_client = KrakenSpotHttpClient::with_credentials(
1755        api_key.into_inner(),
1756        api_secret.into_inner(),
1757        config.environment,
1758        Some(config.http_base_url()),
1759        config.timeout_secs,
1760        None,
1761        None,
1762        None,
1763        config
1764            .proxy_url
1765            .as_ref()
1766            .map(|value| value.expose_secret().to_owned()),
1767        config
1768            .max_requests_per_second
1769            .unwrap_or(KRAKEN_SPOT_DEFAULT_RATE_LIMIT_PER_SECOND),
1770    )
1771    .map_err(|e| {
1772        KrakenWsError::AuthenticationError(format!("Failed to create HTTP client: {e}"))
1773    })?;
1774
1775    let ws_token = http_client.get_websockets_token().await.map_err(|e| {
1776        KrakenWsError::AuthenticationError(format!("Failed to get WebSocket token: {e}"))
1777    })?;
1778
1779    log::debug!(
1780        "WebSocket authentication token refreshed: token_length={}, expires={}",
1781        ws_token.token.expose_secret().len(),
1782        ws_token.expires
1783    );
1784
1785    Ok(ws_token.into_token())
1786}
1787
1788fn update_auth_token_in_payload(
1789    payload: &str,
1790    new_token: &str,
1791) -> Result<SecretString, KrakenWsError> {
1792    let mut value: serde_json::Value =
1793        serde_json::from_str(payload).map_err(|e| KrakenWsError::JsonError(e.to_string()))?;
1794
1795    if let Some(params) = value.get_mut("params") {
1796        params["token"] = serde_json::Value::String(new_token.to_string());
1797    }
1798
1799    serde_json::to_string(&value)
1800        .map(SecretString::from)
1801        .map_err(|e| KrakenWsError::JsonError(e.to_string()))
1802}
1803
1804#[inline]
1805fn to_ws_v2_symbol(symbol: Ustr) -> Ustr {
1806    Ustr::from(&normalize_spot_symbol(symbol.as_str()))
1807}
1808
1809fn bar_type_to_ws_interval(bar_type: BarType) -> Result<u32, KrakenWsError> {
1810    const VALID_INTERVALS: [u32; 9] = [1, 5, 15, 30, 60, 240, 1440, 10080, 21600];
1811
1812    let spec = bar_type.spec();
1813    let step = spec.step.get() as u32;
1814
1815    let base_minutes = match spec.aggregation {
1816        BarAggregation::Minute => 1,
1817        BarAggregation::Hour => 60,
1818        BarAggregation::Day => 1440,
1819        BarAggregation::Week => 10080,
1820        other => {
1821            return Err(KrakenWsError::SubscriptionError(format!(
1822                "Unsupported bar aggregation for Kraken OHLC streaming: {other:?}"
1823            )));
1824        }
1825    };
1826
1827    let interval = base_minutes * step;
1828
1829    if !VALID_INTERVALS.contains(&interval) {
1830        return Err(KrakenWsError::SubscriptionError(format!(
1831            "Invalid bar interval {interval} minutes for Kraken OHLC streaming. \
1832             Supported intervals: 1, 5, 15, 30, 60, 240, 1440, 10080, 21600"
1833        )));
1834    }
1835
1836    Ok(interval)
1837}
1838
1839#[cfg(test)]
1840mod tests {
1841    use std::sync::{Arc, atomic::Ordering};
1842
1843    use log::{Level, LevelFilter, Log, Metadata, Record};
1844    use parking_lot::Mutex;
1845    use rstest::rstest;
1846    use tokio_util::sync::CancellationToken;
1847
1848    use super::*;
1849    use crate::config::KrakenDataClientConfig;
1850
1851    #[rstest]
1852    #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
1853    #[case(
1854        "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
1855        "wss://stream.example.com:9443"
1856    )]
1857    fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
1858        let client = KrakenSpotWebSocketClient::new(
1859            KrakenDataClientConfig {
1860                ws_public_url: Some(url.to_string()),
1861                ..Default::default()
1862            },
1863            CancellationToken::new(),
1864            None,
1865        );
1866
1867        let debug = format!("{client:?}");
1868        assert!(debug.contains(&format!("url: {origin:?}")));
1869        assert!(!debug.contains("SYN-"));
1870        assert_eq!(client.url.expose(), url);
1871    }
1872
1873    #[rstest]
1874    #[tokio::test]
1875    async fn test_debug_redacts_auth_state_and_proxy_url() {
1876        let client = KrakenSpotWebSocketClient::new(
1877            KrakenDataClientConfig::default(),
1878            CancellationToken::new(),
1879            Some("http://user:proxy-secret@localhost".to_string()),
1880        );
1881        *client.auth_token.write().await = Some(SecretString::from("auth-token".to_string()));
1882
1883        let debug = format!("{client:?}");
1884
1885        assert!(debug.contains(REDACTED));
1886        assert!(!debug.contains("proxy-secret"));
1887        assert!(!debug.contains("auth-token"));
1888    }
1889
1890    const SECRET_MARKER: &str = "OUTBOUND_SECRET_MARKER";
1891
1892    struct OutboundLogCapture {
1893        messages: Mutex<Vec<String>>,
1894    }
1895
1896    static OUTBOUND_LOG_CAPTURE: OutboundLogCapture = OutboundLogCapture {
1897        messages: Mutex::new(Vec::new()),
1898    };
1899
1900    impl OutboundLogCapture {
1901        fn clear(&self) {
1902            self.messages.lock().clear();
1903        }
1904
1905        fn messages(&self) -> Vec<String> {
1906            self.messages.lock().clone()
1907        }
1908    }
1909
1910    impl Log for OutboundLogCapture {
1911        fn enabled(&self, metadata: &Metadata<'_>) -> bool {
1912            metadata.level() == Level::Trace
1913                && metadata.target() == "nautilus_kraken::websocket::spot_v2::client"
1914        }
1915
1916        fn log(&self, record: &Record<'_>) {
1917            if self.enabled(record.metadata()) {
1918                let message = record.args().to_string();
1919                if message.starts_with("Sending WebSocket request") {
1920                    self.messages.lock().push(message);
1921                }
1922            }
1923        }
1924
1925        fn flush(&self) {}
1926    }
1927
1928    /// A reconnect replay reuses the original request ids, so a `book` subscribe left unanswered
1929    /// when the socket dropped is retired as the replay opens: a replay's answer matches no
1930    /// request.
1931    #[rstest]
1932    fn test_a_reconnect_replay_retires_the_unanswered_book_requests() {
1933        let client = test_client_without_credentials();
1934        client.record_book_request(7, Ustr::from("BTC/USD"));
1935        assert_eq!(client.book_requests.lock().len(), 1);
1936
1937        begin_resubscribe(
1938            &client.subscriptions,
1939            &client.book_requests,
1940            &client.l2_depths,
1941        );
1942
1943        assert!(client.book_requests.lock().is_empty());
1944    }
1945
1946    #[rstest]
1947    fn test_req_id_counter_is_shared_arc_and_monotonic() {
1948        let cfg = KrakenDataClientConfig::default();
1949        let client = KrakenSpotWebSocketClient::new(cfg, CancellationToken::new(), None);
1950        let counter = client.req_id_counter();
1951        let a = counter.fetch_add(1, Ordering::Relaxed);
1952        let b = counter.fetch_add(1, Ordering::Relaxed);
1953        assert!(b > a);
1954        #[allow(clippy::redundant_clone)]
1955        let cloned = client.clone();
1956        let cloned_counter = cloned.req_id_counter();
1957        assert!(Arc::ptr_eq(&counter, &cloned_counter));
1958    }
1959
1960    #[rstest]
1961    #[case("XBT/EUR", "BTC/EUR")]
1962    #[case("XBT/USD", "BTC/USD")]
1963    #[case("XBT/USDT", "BTC/USDT")]
1964    #[case("ETH/USD", "ETH/USD")]
1965    #[case("ETH/XBT", "ETH/BTC")]
1966    #[case("SOL/XBT", "SOL/BTC")]
1967    #[case("SOL/USD", "SOL/USD")]
1968    #[case("BTC/USD", "BTC/USD")]
1969    #[case("ETH/BTC", "ETH/BTC")]
1970    #[case("XDG/USD", "DOGE/USD")]
1971    #[case("XDG/EUR", "DOGE/EUR")]
1972    fn test_to_kraken_ws_v2_symbol(#[case] input: &str, #[case] expected: &str) {
1973        let symbol = Ustr::from(input);
1974        let result = to_ws_v2_symbol(symbol);
1975        assert_eq!(result, expected);
1976    }
1977
1978    fn test_client_without_credentials() -> KrakenSpotWebSocketClient {
1979        KrakenSpotWebSocketClient::new(
1980            KrakenDataClientConfig::default(),
1981            CancellationToken::new(),
1982            None,
1983        )
1984    }
1985
1986    #[rstest]
1987    #[tokio::test]
1988    async fn test_outbound_logs_omit_payload_bodies() {
1989        log::set_logger(&OUTBOUND_LOG_CAPTURE).expect("test logger already installed");
1990        log::set_max_level(LevelFilter::Trace);
1991
1992        let client = test_client_without_credentials();
1993        let request = KrakenWsRequest {
1994            method: KrakenWsMethod::Subscribe,
1995            params: Some(KrakenWsParams::Channel(KrakenWsChannelParams {
1996                channel: KrakenWsChannel::Executions,
1997                symbol: None,
1998                snapshot: None,
1999                depth: None,
2000                interval: None,
2001                event_trigger: None,
2002                token: Some(SecretString::from(SECRET_MARKER)),
2003                snap_orders: Some(true),
2004                snap_trades: Some(false),
2005            })),
2006            req_id: Some(426),
2007        };
2008        let payload_len = serde_json::to_string(&request).unwrap().len();
2009        OUTBOUND_LOG_CAPTURE.clear();
2010
2011        let error = client.send_command(&request).await.unwrap_err();
2012        let messages = OUTBOUND_LOG_CAPTURE.messages();
2013
2014        assert!(matches!(error, KrakenWsError::ConnectionError(_)));
2015        assert!(
2016            messages
2017                .iter()
2018                .all(|message| !message.contains(SECRET_MARKER)),
2019            "outbound logs exposed the secret marker: {messages:?}"
2020        );
2021        assert!(
2022            messages.iter().any(|message| {
2023                message
2024                    == &format!("Sending WebSocket request: method=Subscribe ({payload_len} bytes)")
2025            }),
2026            "subscribe metadata missing or inaccurate: {messages:?}"
2027        );
2028    }
2029
2030    #[rstest]
2031    #[tokio::test]
2032    async fn test_authenticate_without_credentials_errors() {
2033        let client = test_client_without_credentials();
2034
2035        let err = client.authenticate().await.expect_err("should fail");
2036        assert!(
2037            matches!(err, KrakenWsError::AuthenticationError(ref msg) if msg.contains("API credentials required")),
2038            "unexpected error: {err:?}"
2039        );
2040        assert!(!client.is_authenticated());
2041    }
2042
2043    #[rstest]
2044    #[tokio::test]
2045    async fn test_wait_until_authenticated_times_out() {
2046        let client = test_client_without_credentials();
2047
2048        let err = client
2049            .wait_until_authenticated(0.05)
2050            .await
2051            .expect_err("should time out");
2052        assert!(matches!(err, KrakenWsError::AuthenticationError(_)));
2053    }
2054
2055    #[rstest]
2056    #[tokio::test]
2057    async fn test_wait_until_authenticated_resolves_after_succeed() {
2058        let client = test_client_without_credentials();
2059
2060        let tracker = client.auth_tracker.clone();
2061        let _rx = tracker.begin();
2062
2063        tokio::spawn(async move {
2064            tokio::time::sleep(tokio::time::Duration::from_millis(20)).await;
2065            tracker.succeed();
2066        });
2067
2068        client
2069            .wait_until_authenticated(1.0)
2070            .await
2071            .expect("should resolve once tracker succeeds");
2072        assert!(client.is_authenticated());
2073    }
2074
2075    #[rstest]
2076    #[tokio::test]
2077    async fn test_is_authenticated_flips_on_fail() {
2078        let client = test_client_without_credentials();
2079
2080        let _rx = client.auth_tracker.begin();
2081        client.auth_tracker.succeed();
2082        assert!(client.is_authenticated());
2083
2084        client.auth_tracker.fail("test failure");
2085        assert!(!client.is_authenticated());
2086    }
2087
2088    #[rstest]
2089    fn test_l3_factory_uses_ws_l3_url() {
2090        let cfg = KrakenDataClientConfig::default();
2091        let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2092        assert_eq!(client.url(), "wss://ws-l3.kraken.com/v2");
2093    }
2094
2095    #[rstest]
2096    fn test_l3_factory_respects_override() {
2097        let cfg = KrakenDataClientConfig {
2098            ws_l3_url: Some("wss://override.example/v2".to_string()),
2099            ..Default::default()
2100        };
2101        let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2102        assert_eq!(client.url(), "wss://override.example/v2");
2103    }
2104
2105    #[rstest]
2106    #[tokio::test]
2107    async fn test_subscribe_book_l3_without_auth_errors_and_leaves_clean_state() {
2108        let cfg = KrakenDataClientConfig::default();
2109        let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2110
2111        let err = client
2112            .subscribe_book_l3(Ustr::from("BTC/USD"), 1000)
2113            .await
2114            .expect_err("should fail without auth token");
2115
2116        assert!(matches!(err, KrakenWsError::AuthenticationError(_)));
2117        assert!(
2118            client.subscriptions.is_empty(),
2119            "no state must leak on auth failure"
2120        );
2121    }
2122
2123    #[rstest]
2124    #[tokio::test]
2125    async fn test_subscribe_book_l3_invalid_depth_errors() {
2126        let cfg = KrakenDataClientConfig::default();
2127        let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2128
2129        let err = client
2130            .subscribe_book_l3(Ustr::from("BTC/USD"), 50)
2131            .await
2132            .expect_err("should fail on invalid depth");
2133
2134        assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2135        assert!(!client.subscriptions_contains("level3:BTC/USD"));
2136    }
2137
2138    #[rstest]
2139    #[tokio::test]
2140    async fn test_subscribe_book_invalid_depth_errors() {
2141        let client = test_client_without_credentials();
2142        let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2143
2144        let err = client
2145            .subscribe_book(instrument_id, Some(50))
2146            .await
2147            .expect_err("should fail on invalid depth");
2148
2149        assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2150        assert!(!client.subscriptions_contains("book:BTC/USD"));
2151        assert_eq!(client.l2_depths.get("BTC/USD"), None);
2152    }
2153
2154    #[rstest]
2155    #[tokio::test]
2156    async fn test_subscribe_book_defaults_to_depth_10_and_stores_state() {
2157        let client = test_client_without_credentials();
2158        let key = "book:BTC/USD";
2159        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2160        *client.cmd_tx.write().await = cmd_tx;
2161
2162        client
2163            .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), None)
2164            .await
2165            .expect("subscribe should succeed");
2166
2167        let cmd = cmd_rx.try_recv().expect("expected subscribe command");
2168        let SpotHandlerCommand::Subscribe { payload } = cmd else {
2169            panic!("expected subscribe command");
2170        };
2171        assert_book_subscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2172
2173        assert_eq!(client.subscriptions.get_reference_count(key), 1);
2174        assert!(client.subscriptions_contains(key));
2175        assert_eq!(client.l2_depths.get("BTC/USD"), Some(10));
2176        assert_eq!(
2177            client.subscription_payloads.read().await.get(key),
2178            Some(&payload)
2179        );
2180    }
2181
2182    #[rstest]
2183    #[tokio::test]
2184    async fn test_subscribe_book_send_failure_rolls_back_l2_state() {
2185        let client = test_client_without_credentials();
2186        let key = "book:BTC/USD";
2187
2188        let err = client
2189            .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(10))
2190            .await
2191            .expect_err("should fail when command receiver is closed");
2192
2193        assert!(matches!(err, KrakenWsError::ConnectionError(_)));
2194        assert_eq!(client.subscriptions.get_reference_count(key), 0);
2195        assert!(!client.subscriptions_contains(key));
2196        assert!(client.subscriptions.pending_subscribe_topics().is_empty());
2197        assert!(client.subscriptions.pending_unsubscribe_topics().is_empty());
2198        assert_eq!(client.l2_depths.get("BTC/USD"), None);
2199        assert!(!client.subscription_payloads.read().await.contains_key(key));
2200    }
2201
2202    #[rstest]
2203    fn test_subscribe_book_l3_refcount_idempotent() {
2204        let cfg = KrakenDataClientConfig::default();
2205        let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2206        let key = "level3:BTC/USD";
2207
2208        assert!(client.subscriptions.add_reference(key));
2209        assert!(!client.subscriptions.add_reference(key));
2210        assert!(!client.subscriptions.remove_reference(key));
2211        assert!(client.subscriptions.remove_reference(key));
2212    }
2213
2214    #[rstest]
2215    #[tokio::test]
2216    async fn test_subscribe_book_l3_rejects_depth_mismatch() {
2217        let cfg = KrakenDataClientConfig::default();
2218        let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2219
2220        let key = "level3:BTC/USD";
2221        client.subscriptions.add_reference(key);
2222        client.subscriptions.mark_subscribe(key);
2223        client.subscriptions.confirm_subscribe(key);
2224        client.l3_depths.lock().insert("BTC/USD".to_string(), 1000);
2225
2226        *client.auth_token.write().await = Some(SecretString::from("test-token".to_string()));
2227
2228        let err = client
2229            .subscribe_book_l3(Ustr::from("BTC/USD"), 10)
2230            .await
2231            .expect_err("should reject depth mismatch");
2232        assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2233
2234        assert!(client.subscriptions.remove_reference(key));
2235    }
2236
2237    #[rstest]
2238    #[tokio::test]
2239    async fn test_subscribe_book_rejects_depth_mismatch() {
2240        let client = test_client_without_credentials();
2241        let key = "book:BTC/USD";
2242        client.subscriptions.add_reference(key);
2243        client.subscriptions.mark_subscribe(key);
2244        client.subscriptions.confirm_subscribe(key);
2245        client.l2_depths.insert("BTC/USD", 10, 1);
2246
2247        let err = client
2248            .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(25))
2249            .await
2250            .expect_err("should reject depth mismatch");
2251        assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2252
2253        assert!(client.subscriptions.remove_reference(key));
2254    }
2255
2256    #[rstest]
2257    #[tokio::test]
2258    async fn test_unsubscribe_book_removes_l2_depth_on_last_reference() {
2259        let client = test_client_without_credentials();
2260        let key = "book:BTC/USD";
2261        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2262        *client.cmd_tx.write().await = cmd_tx;
2263
2264        client.subscriptions.add_reference(key);
2265        client.subscriptions.mark_subscribe(key);
2266        client.subscriptions.confirm_subscribe(key);
2267        client.l2_depths.insert("BTC/USD", 10, 1);
2268        client
2269            .subscription_payloads
2270            .write()
2271            .await
2272            .insert(key.to_string(), SecretString::from("payload".to_string()));
2273
2274        client
2275            .unsubscribe_book(InstrumentId::from("BTC/USD.KRAKEN"))
2276            .await
2277            .expect("unsubscribe should succeed");
2278
2279        assert_eq!(client.subscriptions.get_reference_count(key), 0);
2280        assert!(!client.subscriptions_contains(key));
2281        assert_eq!(client.l2_depths.get("BTC/USD"), None);
2282        assert!(!client.subscription_payloads.read().await.contains_key(key));
2283
2284        let cmd = cmd_rx.try_recv().expect("expected unsubscribe command");
2285        let SpotHandlerCommand::Unsubscribe { payload } = cmd else {
2286            panic!("expected unsubscribe command");
2287        };
2288        assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2289    }
2290
2291    #[rstest]
2292    #[tokio::test]
2293    async fn test_generic_subscribe_rejects_level3() {
2294        let cfg = KrakenDataClientConfig::default();
2295        let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2296
2297        let err = client
2298            .subscribe(
2299                KrakenWsChannel::Level3,
2300                vec![Ustr::from("BTC/USD")],
2301                Some(1000),
2302            )
2303            .await
2304            .expect_err("generic subscribe must reject Level3");
2305        assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2306
2307        let err = client
2308            .unsubscribe(KrakenWsChannel::Level3, vec![Ustr::from("BTC/USD")])
2309            .await
2310            .expect_err("generic unsubscribe must reject Level3");
2311        assert!(matches!(err, KrakenWsError::InvalidMessage(_)));
2312    }
2313
2314    #[rstest]
2315    fn test_update_auth_token_in_payload_for_level3() {
2316        let original = r#"{"method":"subscribe","params":{"channel":"level3","symbol":["BTC/USD"],"depth":1000,"snapshot":true,"token":"OLD"},"req_id":1}"#;
2317        let rewritten = update_auth_token_in_payload(original, "NEW").unwrap();
2318        assert!(rewritten.expose_secret().contains(r#""token":"NEW""#));
2319        assert!(!rewritten.expose_secret().contains(r#""token":"OLD""#));
2320    }
2321
2322    #[rstest]
2323    fn test_l3_depths_shared_between_subscribe_and_handle() {
2324        let cfg = KrakenDataClientConfig::default();
2325        let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2326
2327        let handle = client.l3_depths_handle();
2328        client.l3_depths.lock().insert("BTC/USD".to_string(), 100);
2329
2330        assert_eq!(handle.lock().get("BTC/USD").copied(), Some(100));
2331    }
2332
2333    #[rstest]
2334    fn test_l2_depths_shared_between_subscribe_and_handle() {
2335        let client = test_client_without_credentials();
2336
2337        let handle = client.l2_depths_handle();
2338        client.l2_depths.insert("BTC/USD", 10, 1);
2339
2340        assert_eq!(handle.get("BTC/USD"), Some(10));
2341    }
2342
2343    /// Every `book` subscribe the client sends is on record under its request id with the pair,
2344    /// so the venue's answer can be matched to it, and is the symbol's latest request.
2345    #[rstest]
2346    #[tokio::test]
2347    async fn test_subscribe_and_resync_book_record_their_requests() {
2348        let client = test_client_without_credentials();
2349        let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2350        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2351        *client.cmd_tx.write().await = cmd_tx;
2352        let handle = client.book_requests_handle();
2353
2354        client
2355            .subscribe_book(instrument_id, Some(10))
2356            .await
2357            .unwrap();
2358        let subscription = client.l2_depths.subscription("BTC/USD").unwrap();
2359        let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2360        else {
2361            panic!("expected a subscribe command");
2362        };
2363        let subscribe_req_id = payload_req_id(payload.expose_secret());
2364
2365        assert_eq!(
2366            handle.lock().get(&subscribe_req_id),
2367            Some(&L2BookRequest {
2368                symbol: Ustr::from("BTC/USD"),
2369            })
2370        );
2371        assert_eq!(subscription.latest_request, subscribe_req_id);
2372
2373        client
2374            .resync_book(
2375                instrument_id,
2376                subscription.generation,
2377                subscription.snapshot_epoch,
2378            )
2379            .await
2380            .unwrap();
2381        let SpotHandlerCommand::Unsubscribe { payload } = cmd_rx.try_recv().expect("unsubscribe")
2382        else {
2383            panic!("expected an unsubscribe command");
2384        };
2385        assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2386        let unsubscribe_req_id = payload_req_id(payload.expose_secret());
2387        let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2388        else {
2389            panic!("expected a subscribe command");
2390        };
2391        let resync_req_id = payload_req_id(payload.expose_secret());
2392
2393        let requests = handle.lock();
2394        assert_eq!(
2395            requests.get(&resync_req_id),
2396            Some(&L2BookRequest {
2397                symbol: Ustr::from("BTC/USD"),
2398            })
2399        );
2400        assert_eq!(
2401            client
2402                .l2_depths
2403                .subscription("BTC/USD")
2404                .unwrap()
2405                .latest_request,
2406            resync_req_id
2407        );
2408        assert!(
2409            !requests.contains_key(&unsubscribe_req_id),
2410            "only subscribes are on record"
2411        );
2412        assert_eq!(requests.len(), 2);
2413    }
2414
2415    /// A disconnect clears the record of in-flight subscribes along with the subscriptions: the
2416    /// venue answers none of them on a new connection.
2417    #[rstest]
2418    #[tokio::test]
2419    async fn test_disconnect_clears_the_book_requests() {
2420        let client = test_client_without_credentials();
2421        let (cmd_tx, _cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2422        *client.cmd_tx.write().await = cmd_tx;
2423        client
2424            .subscribe_book(InstrumentId::from("BTC/USD.KRAKEN"), Some(10))
2425            .await
2426            .unwrap();
2427        assert_eq!(client.book_requests.lock().len(), 1);
2428
2429        let _ = client.disconnect_locked().await;
2430
2431        assert!(client.book_requests.lock().is_empty());
2432    }
2433
2434    fn payload_req_id(payload: &str) -> u64 {
2435        let value: serde_json::Value = serde_json::from_str(payload).unwrap();
2436        value["req_id"].as_u64().expect("a req_id")
2437    }
2438
2439    /// A recovery whose snapshot the data client has accepted since it was issued sends nothing,
2440    /// leaving the recovered stream subscribed; a recovery carrying the epoch that snapshot moved
2441    /// to, as a mismatch on the snapshot issues, resubscribes.
2442    #[rstest]
2443    #[tokio::test]
2444    async fn test_resync_book_skips_once_a_snapshot_is_accepted() {
2445        let client = test_client_without_credentials();
2446        let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2447        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2448        *client.cmd_tx.write().await = cmd_tx;
2449        client
2450            .subscribe_book(instrument_id, Some(10))
2451            .await
2452            .unwrap();
2453        while cmd_rx.try_recv().is_ok() {}
2454        let issued = client.l2_depths.subscription("BTC/USD").unwrap();
2455
2456        let accepted = client
2457            .l2_depths
2458            .accept_snapshot("BTC/USD", issued.latest_request)
2459            .expect("the live stream's snapshot is accepted");
2460        client
2461            .resync_book(instrument_id, issued.generation, issued.snapshot_epoch)
2462            .await
2463            .unwrap();
2464
2465        assert!(
2466            cmd_rx.try_recv().is_err(),
2467            "a served recovery must send nothing"
2468        );
2469        assert_eq!(
2470            client
2471                .l2_depths
2472                .subscription("BTC/USD")
2473                .unwrap()
2474                .latest_request,
2475            issued.latest_request
2476        );
2477
2478        client
2479            .resync_book(instrument_id, accepted.generation, accepted.snapshot_epoch)
2480            .await
2481            .unwrap();
2482        assert_book_recovery_commands(&mut cmd_rx, 10);
2483    }
2484
2485    /// A recovery issued before a reconnect sends nothing once the replay has begun, so it does
2486    /// not cycle the replayed stream; one issued after it resubscribes.
2487    #[rstest]
2488    #[tokio::test]
2489    async fn test_a_recovery_issued_before_a_reconnect_sends_nothing() {
2490        let client = test_client_without_credentials();
2491        let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2492        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2493        *client.cmd_tx.write().await = cmd_tx;
2494        client
2495            .subscribe_book(instrument_id, Some(10))
2496            .await
2497            .unwrap();
2498        while cmd_rx.try_recv().is_ok() {}
2499        let issued = client.l2_depths.subscription("BTC/USD").unwrap();
2500
2501        begin_resubscribe(
2502            &client.subscriptions,
2503            &client.book_requests,
2504            &client.l2_depths,
2505        );
2506        client
2507            .resync_book(instrument_id, issued.generation, issued.snapshot_epoch)
2508            .await
2509            .unwrap();
2510
2511        assert!(
2512            cmd_rx.try_recv().is_err(),
2513            "a recovery from before the reconnect must send nothing"
2514        );
2515
2516        let replayed = client.l2_depths.subscription("BTC/USD").unwrap();
2517        client
2518            .resync_book(instrument_id, replayed.generation, replayed.snapshot_epoch)
2519            .await
2520            .unwrap();
2521        assert_book_recovery_commands(&mut cmd_rx, 10);
2522    }
2523
2524    /// A recovery records its resubscribe as the symbol's latest request before sending and keeps
2525    /// it when a send fails, so the data client sees the request unconfirmed and its watchdog asks
2526    /// again rather than accepting the old stream as current.
2527    #[rstest]
2528    #[tokio::test]
2529    async fn test_resync_book_keeps_its_latest_request_when_a_send_fails() {
2530        let client = test_client_without_credentials();
2531        let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2532        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2533        *client.cmd_tx.write().await = cmd_tx;
2534        client
2535            .subscribe_book(instrument_id, Some(10))
2536            .await
2537            .unwrap();
2538        let issued = client.l2_depths.subscription("BTC/USD").unwrap();
2539        drop(cmd_rx);
2540
2541        client
2542            .resync_book(instrument_id, issued.generation, issued.snapshot_epoch)
2543            .await
2544            .expect_err("the command channel is closed");
2545
2546        let latest = client
2547            .l2_depths
2548            .subscription("BTC/USD")
2549            .unwrap()
2550            .latest_request;
2551        assert_ne!(latest, issued.latest_request);
2552        assert!(
2553            !client.book_requests.lock().contains_key(&latest),
2554            "an unsent subscribe has no answer on record"
2555        );
2556    }
2557
2558    #[rstest]
2559    fn test_resync_book_l3_does_not_touch_refcount() {
2560        let cfg = KrakenDataClientConfig::default();
2561        let client = KrakenSpotWebSocketClient::l3(cfg, CancellationToken::new(), None);
2562        let key = "level3:BTC/USD";
2563
2564        assert!(client.subscriptions.add_reference(key));
2565        assert!(!client.subscriptions.add_reference(key));
2566        client.subscriptions.mark_subscribe(key);
2567        client.subscriptions.confirm_subscribe(key);
2568
2569        assert!(client.subscriptions_contains(key));
2570    }
2571
2572    /// A recovery queued for a subscription the user has since replaced leaves the replacement
2573    /// alone: it sends nothing and the recorded depth is the replacement's.
2574    #[rstest]
2575    #[tokio::test]
2576    async fn test_resync_book_leaves_a_replacement_subscription_alone() {
2577        let client = test_client_without_credentials();
2578        let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2579        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2580        *client.cmd_tx.write().await = cmd_tx;
2581
2582        client
2583            .subscribe_book(instrument_id, Some(10))
2584            .await
2585            .unwrap();
2586        let retired = client.l2_depths.subscription("BTC/USD").unwrap().generation;
2587        client.unsubscribe_book(instrument_id).await.unwrap();
2588        client
2589            .subscribe_book(instrument_id, Some(100))
2590            .await
2591            .unwrap();
2592
2593        while cmd_rx.try_recv().is_ok() {}
2594
2595        client.resync_book(instrument_id, retired, 0).await.unwrap();
2596
2597        assert!(
2598            cmd_rx.try_recv().is_err(),
2599            "a retired subscription's recovery must send nothing"
2600        );
2601        assert_eq!(client.l2_depths.get("BTC/USD"), Some(100));
2602
2603        // Control: the live subscription's own recovery resubscribes at its depth.
2604        let live = client.l2_depths.subscription("BTC/USD").unwrap();
2605        client
2606            .resync_book(instrument_id, live.generation, live.snapshot_epoch)
2607            .await
2608            .unwrap();
2609
2610        let SpotHandlerCommand::Unsubscribe { payload } = cmd_rx.try_recv().expect("unsubscribe")
2611        else {
2612            panic!("expected an unsubscribe command");
2613        };
2614        assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", 100);
2615        let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2616        else {
2617            panic!("expected a subscribe command");
2618        };
2619        assert_book_subscribe_payload(payload.expose_secret(), "BTC/USD", 100);
2620        assert_eq!(client.l2_depths.get("BTC/USD"), Some(100));
2621    }
2622
2623    /// Spawns `command` and yields until it has finished or parked on a lock, so any command it
2624    /// could send without waiting is already in the queue.
2625    async fn spawn_parked<F>(command: F) -> tokio::task::JoinHandle<Result<(), KrakenWsError>>
2626    where
2627        F: Future<Output = Result<(), KrakenWsError>> + Send + 'static,
2628    {
2629        let handle = tokio::spawn(command);
2630
2631        for _ in 0..10 {
2632            tokio::task::yield_now().await;
2633        }
2634
2635        handle
2636    }
2637
2638    fn assert_book_recovery_commands(
2639        cmd_rx: &mut tokio::sync::mpsc::UnboundedReceiver<SpotHandlerCommand>,
2640        depth: u32,
2641    ) {
2642        let SpotHandlerCommand::Unsubscribe { payload } = cmd_rx.try_recv().expect("unsubscribe")
2643        else {
2644            panic!("expected an unsubscribe command");
2645        };
2646        assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", depth);
2647        let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2648        else {
2649            panic!("expected a subscribe command");
2650        };
2651        assert_book_subscribe_payload(payload.expose_secret(), "BTC/USD", depth);
2652        let value: serde_json::Value = serde_json::from_str(payload.expose_secret()).unwrap();
2653        assert_eq!(value["params"]["snapshot"], serde_json::json!(true));
2654        assert!(cmd_rx.try_recv().is_err(), "nothing else is sent");
2655    }
2656
2657    /// A recovery waits for the book command lock, so its check of the live subscription and its
2658    /// two commands cannot interleave with a subscription change holding the lock.
2659    #[rstest]
2660    #[tokio::test]
2661    async fn test_resync_book_waits_for_the_book_command_lock() {
2662        let client = test_client_without_credentials();
2663        let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2664        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2665        *client.cmd_tx.write().await = cmd_tx;
2666        client
2667            .subscribe_book(instrument_id, Some(25))
2668            .await
2669            .unwrap();
2670        while cmd_rx.try_recv().is_ok() {}
2671        let live = client.l2_depths.subscription("BTC/USD").unwrap();
2672
2673        let guard = client.l2_book_commands.lock().await;
2674        let task_client = client.clone();
2675        let recovery = spawn_parked(async move {
2676            task_client
2677                .resync_book(instrument_id, live.generation, live.snapshot_epoch)
2678                .await
2679        })
2680        .await;
2681        assert!(
2682            cmd_rx.try_recv().is_err(),
2683            "no command while a subscription change holds the lock"
2684        );
2685
2686        drop(guard);
2687        recovery.await.unwrap().unwrap();
2688        assert_book_recovery_commands(&mut cmd_rx, 25);
2689    }
2690
2691    /// A subscription change waits for the book command lock, so it cannot land between a
2692    /// recovery's check and its commands.
2693    #[rstest]
2694    #[tokio::test]
2695    async fn test_subscribe_and_unsubscribe_book_wait_for_the_book_command_lock() {
2696        let client = test_client_without_credentials();
2697        let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2698        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2699        *client.cmd_tx.write().await = cmd_tx;
2700
2701        let guard = client.l2_book_commands.lock().await;
2702        let task_client = client.clone();
2703        let subscribe =
2704            spawn_parked(async move { task_client.subscribe_book(instrument_id, Some(10)).await })
2705                .await;
2706        assert!(
2707            cmd_rx.try_recv().is_err(),
2708            "no subscribe while the lock is held"
2709        );
2710
2711        drop(guard);
2712        subscribe.await.unwrap().unwrap();
2713        let SpotHandlerCommand::Subscribe { payload } = cmd_rx.try_recv().expect("subscribe")
2714        else {
2715            panic!("expected a subscribe command");
2716        };
2717        assert_book_subscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2718        assert!(cmd_rx.try_recv().is_err());
2719
2720        let guard = client.l2_book_commands.lock().await;
2721        let task_client = client.clone();
2722        let unsubscribe =
2723            spawn_parked(async move { task_client.unsubscribe_book(instrument_id).await }).await;
2724        assert!(
2725            cmd_rx.try_recv().is_err(),
2726            "no unsubscribe while the lock is held"
2727        );
2728
2729        drop(guard);
2730        unsubscribe.await.unwrap().unwrap();
2731        let SpotHandlerCommand::Unsubscribe { payload } = cmd_rx.try_recv().expect("unsubscribe")
2732        else {
2733            panic!("expected an unsubscribe command");
2734        };
2735        assert_book_unsubscribe_payload(payload.expose_secret(), "BTC/USD", 10);
2736        assert!(cmd_rx.try_recv().is_err());
2737        assert_eq!(client.l2_depths.get("BTC/USD"), None);
2738    }
2739
2740    /// Once admitted, a recovery sends both commands: a stale unsubscribe with no subscribe behind
2741    /// it would leave the stream dead, so the depth seen at admission is the one resubscribed.
2742    #[rstest]
2743    #[tokio::test]
2744    async fn test_resync_book_sends_both_commands_once_admitted() {
2745        let client = test_client_without_credentials();
2746        let instrument_id = InstrumentId::from("BTC/USD.KRAKEN");
2747        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
2748        *client.cmd_tx.write().await = cmd_tx;
2749        client
2750            .subscribe_book(instrument_id, Some(10))
2751            .await
2752            .unwrap();
2753        while cmd_rx.try_recv().is_ok() {}
2754        let live = client.l2_depths.subscription("BTC/USD").unwrap();
2755
2756        // Parks `send_command` at admission: the recovery has passed its check and holds the lock.
2757        let admission = client.cmd_tx.write().await;
2758        let task_client = client.clone();
2759        let recovery = spawn_parked(async move {
2760            task_client
2761                .resync_book(instrument_id, live.generation, live.snapshot_epoch)
2762                .await
2763        })
2764        .await;
2765        assert!(cmd_rx.try_recv().is_err());
2766
2767        // Models a depth change observed after the check and before the commands go out.
2768        client.l2_depths.insert("BTC/USD", 25, 99);
2769        drop(admission);
2770        recovery.await.unwrap().unwrap();
2771
2772        assert_book_recovery_commands(&mut cmd_rx, 10);
2773    }
2774
2775    fn assert_book_subscribe_payload(payload: &str, symbol: &str, depth: u32) {
2776        let value: serde_json::Value =
2777            serde_json::from_str(payload).expect("payload should parse as JSON");
2778
2779        assert_eq!(value["method"], serde_json::json!("subscribe"));
2780        assert_eq!(value["params"]["channel"], serde_json::json!("book"));
2781        assert_eq!(value["params"]["symbol"], serde_json::json!([symbol]));
2782        assert_eq!(value["params"]["depth"], serde_json::json!(depth));
2783    }
2784
2785    /// The unsubscribe names the depth: the venue keys a `book` subscription by symbol and depth.
2786    fn assert_book_unsubscribe_payload(payload: &str, symbol: &str, depth: u32) {
2787        let value: serde_json::Value =
2788            serde_json::from_str(payload).expect("payload should parse as JSON");
2789
2790        assert_eq!(value["method"], serde_json::json!("unsubscribe"));
2791        assert_eq!(value["params"]["channel"], serde_json::json!("book"));
2792        assert_eq!(value["params"]["symbol"], serde_json::json!([symbol]));
2793        assert_eq!(value["params"]["depth"], serde_json::json!(depth));
2794    }
2795}