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