Skip to main content

nautilus_coinbase/websocket/
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 Coinbase Advanced Trade API.
17//!
18//! Manages connection lifecycle, JWT-authenticated subscriptions, and dispatches
19//! parsed Nautilus messages through the [`FeedHandler`].
20
21use std::{
22    num::NonZeroU32,
23    str::FromStr,
24    sync::{
25        Arc, LazyLock,
26        atomic::{AtomicBool, AtomicU8, Ordering},
27    },
28    time::Duration,
29};
30
31use arc_swap::ArcSwap;
32#[cfg(test)]
33use nautilus_core::string::secret::REDACTED;
34use nautilus_core::{AtomicMap, string::secret::SecretString};
35use nautilus_live::{
36    SocketControl,
37    task::{TaskJoinOutcome, TaskSlot, finish_task},
38};
39use nautilus_model::{
40    data::BarType,
41    identifiers::{AccountId, InstrumentId},
42    instruments::{Instrument, InstrumentAny},
43};
44use nautilus_network::{
45    ConnectionUrl,
46    http::create_standard_nautilus_headers,
47    mode::ConnectionMode,
48    ratelimiter::quota::Quota,
49    websocket::{
50        SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
51        channel_message_handler,
52    },
53};
54use ustr::Ustr;
55
56use crate::{
57    common::{
58        consts::{
59            RECONNECT_BACKOFF_FACTOR, RECONNECT_BASE_BACKOFF, RECONNECT_JITTER_MS,
60            RECONNECT_MAX_BACKOFF, RECONNECT_TIMEOUT, WS_DISCONNECT_TIMEOUT, WS_HEARTBEAT_SECS,
61        },
62        credential::CoinbaseCredential,
63        enums::CoinbaseWsChannel,
64    },
65    websocket::{
66        handler::{FeedHandler, HandlerCommand, NautilusWsMessage},
67        messages::{CoinbaseWsAction, CoinbaseWsSubscription},
68    },
69};
70
71/// Coinbase WebSocket connection rate limit (8 per second per IP).
72pub static COINBASE_WS_CONNECTION_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
73    Quota::per_second(NonZeroU32::new(8).expect("non-zero")).expect("valid constant")
74});
75
76/// Coinbase WebSocket subscribe/unsubscribe rate limit (8 per second per IP).
77pub static COINBASE_WS_SUBSCRIPTION_QUOTA: LazyLock<Quota> = LazyLock::new(|| {
78    Quota::per_second(NonZeroU32::new(8).expect("non-zero")).expect("valid constant")
79});
80
81/// Rate-limit key for subscribe/unsubscribe operations.
82pub const COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION: &str = "subscription";
83
84/// Pre-interned [`COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION`] slice.
85pub static COINBASE_WS_SUBSCRIPTION_KEYS: LazyLock<[Ustr; 1]> =
86    LazyLock::new(|| [Ustr::from(COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION)]);
87
88/// WebSocket client for Coinbase Advanced Trade market data and user streams.
89///
90/// Manages connection lifecycle, subscription state, and JWT authentication.
91/// Spawns a [`FeedHandler`] task that parses raw messages into Nautilus types.
92#[derive(Debug)]
93pub struct CoinbaseWebSocketClient {
94    url: ConnectionUrl,
95    connection_mode: Arc<ArcSwap<AtomicU8>>,
96    signal: Arc<AtomicBool>,
97    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
98    out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<NautilusWsMessage>>,
99    instruments: Arc<AtomicMap<InstrumentId, InstrumentAny>>,
100    /// Maps a canonical wire `product_id` to the `product_id` the caller
101    /// subscribed or submitted with. Coinbase rewrites aliased products to
102    /// their canonical form on the wire (e.g. `BTC-USDC -> BTC-USD`), so
103    /// inbound messages must be re-keyed to the caller's id before parsing.
104    subscription_aliases: Arc<AtomicMap<Ustr, Ustr>>,
105    bar_types: ahash::AHashMap<String, BarType>,
106    subscriptions: SubscriptionState,
107    credential: Option<CoinbaseCredential>,
108    account_id: Option<AccountId>,
109    task_handle: TaskSlot<()>,
110    shutdown_errors: Vec<String>,
111    transport_backend: TransportBackend,
112    proxy_url: Option<SecretString>,
113    socket_control: Option<SocketControl>,
114}
115
116impl Clone for CoinbaseWebSocketClient {
117    fn clone(&self) -> Self {
118        Self {
119            url: self.url.clone(),
120            connection_mode: Arc::clone(&self.connection_mode),
121            signal: Arc::clone(&self.signal),
122            cmd_tx: Arc::clone(&self.cmd_tx),
123            out_rx: None,
124            instruments: Arc::clone(&self.instruments),
125            subscription_aliases: Arc::clone(&self.subscription_aliases),
126            bar_types: self.bar_types.clone(),
127            subscriptions: self.subscriptions.clone(),
128            credential: self.credential.clone(),
129            account_id: self.account_id,
130            task_handle: TaskSlot::new(),
131            shutdown_errors: Vec::new(),
132            transport_backend: self.transport_backend,
133            proxy_url: self.proxy_url.clone(),
134            socket_control: self.socket_control.clone(),
135        }
136    }
137}
138
139impl CoinbaseWebSocketClient {
140    /// Creates a new [`CoinbaseWebSocketClient`] for public market data.
141    pub fn new(url: &str, transport_backend: TransportBackend, proxy_url: Option<String>) -> Self {
142        let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
143
144        Self {
145            url: url.into(),
146            connection_mode: Arc::new(ArcSwap::from_pointee(AtomicU8::new(
147                ConnectionMode::Closed.as_u8(),
148            ))),
149            signal: Arc::new(AtomicBool::new(false)),
150            cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
151            out_rx: None,
152            instruments: Arc::new(AtomicMap::new()),
153            subscription_aliases: Arc::new(AtomicMap::new()),
154            bar_types: ahash::AHashMap::new(),
155            subscriptions: SubscriptionState::new('|'),
156            credential: None,
157            account_id: None,
158            task_handle: TaskSlot::new(),
159            shutdown_errors: Vec::new(),
160            transport_backend,
161            proxy_url: proxy_url.map(SecretString::from),
162            socket_control: None,
163        }
164    }
165
166    /// Configures socket state reporting and reconnect control.
167    #[must_use]
168    pub fn with_socket_control(mut self, control: SocketControl) -> Self {
169        self.socket_control = Some(control);
170        self
171    }
172
173    /// Creates a new [`CoinbaseWebSocketClient`] with credentials for authenticated channels.
174    pub fn with_credential(
175        url: &str,
176        credential: CoinbaseCredential,
177        transport_backend: TransportBackend,
178        proxy_url: Option<String>,
179    ) -> Self {
180        let mut client = Self::new(url, transport_backend, proxy_url);
181        client.credential = Some(credential);
182        client
183    }
184
185    /// Sets the account ID used when emitting user-channel execution reports.
186    ///
187    /// Propagates to the feed handler when the connection is active so that
188    /// subsequent user events carry the correct account identifier.
189    pub async fn set_account_id(&mut self, account_id: AccountId) {
190        self.account_id = Some(account_id);
191
192        let cmd_tx = self.cmd_tx.read().await;
193        if let Err(e) = cmd_tx.send(HandlerCommand::SetAccountId(account_id)) {
194            log::debug!("Failed to send SetAccountId: {e}");
195        }
196    }
197
198    /// Bulk-populates the instrument cache.
199    ///
200    /// Safe to call before or after [`Self::connect`]. When called before
201    /// connect, instruments are picked up by the initial `InitializeInstruments`
202    /// command the client sends to the handler; when called after, a fresh
203    /// `InitializeInstruments` command is sent to refresh the handler's cache.
204    pub async fn initialize_instruments(&self, instruments: Vec<InstrumentAny>) {
205        for instrument in &instruments {
206            self.instruments.insert(instrument.id(), instrument.clone());
207        }
208
209        let cmd_tx = self.cmd_tx.read().await;
210        if let Err(e) = cmd_tx.send(HandlerCommand::InitializeInstruments(instruments)) {
211            log::debug!("Failed to send InitializeInstruments: {e}");
212        }
213    }
214
215    // Coinbase closes clients that idle without a subscribe inside 5s, and
216    // heartbeats keeps the connection alive when product topics are quiet.
217    // Marking before `resubscribe_all` replays it on every reconnect.
218    fn prime_default_subscriptions(&self) {
219        self.subscriptions
220            .mark_subscribe(CoinbaseWsChannel::Heartbeats.as_ref());
221    }
222
223    /// Establishes the WebSocket connection and spawns the feed handler.
224    pub async fn connect(&mut self) -> anyhow::Result<()> {
225        if self.is_active() || self.is_reconnecting() {
226            log::warn!("WebSocket already connected or reconnecting");
227            return Ok(());
228        }
229
230        if let Some(outcome) = finish_task(
231            &mut self.task_handle,
232            WS_DISCONNECT_TIMEOUT,
233            WS_DISCONNECT_TIMEOUT,
234        )
235        .await
236        {
237            match outcome {
238                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
239                TaskJoinOutcome::Failed(error) => {
240                    anyhow::bail!("Coinbase WebSocket handler failed: {error}");
241                }
242                TaskJoinOutcome::Incomplete => {
243                    anyhow::bail!("Coinbase WebSocket handler did not stop after abort");
244                }
245            }
246        }
247
248        // Clear stop signal from any previous disconnect
249        self.signal.store(false, Ordering::Relaxed);
250
251        let (message_handler, raw_rx) = channel_message_handler();
252        let headers = create_standard_nautilus_headers();
253
254        let cfg = WebSocketConfig {
255            url: self.url.expose().to_owned(),
256            headers,
257            // Coinbase uses TCP control-frame pings for transport keep-alive;
258            // application-layer liveness comes from the heartbeats channel.
259            heartbeat_interval_secs: Some(WS_HEARTBEAT_SECS),
260            heartbeat_payload: None,
261            connect_timeout_ms: Some(RECONNECT_TIMEOUT.as_millis() as u64),
262            reconnect_delay_initial_ms: Some(RECONNECT_BASE_BACKOFF.as_millis() as u64),
263            reconnect_delay_max_ms: Some(RECONNECT_MAX_BACKOFF.as_millis() as u64),
264            reconnect_backoff_factor: Some(RECONNECT_BACKOFF_FACTOR),
265            reconnect_jitter_ms: Some(RECONNECT_JITTER_MS),
266            reconnect_max_attempts: None,
267            heartbeat_timeout_secs: None,
268            idle_timeout_ms: None,
269            writer_capacity: None,
270            backend: self.transport_backend,
271            proxy_url: self
272                .proxy_url
273                .as_ref()
274                .map(|value| value.expose_secret().to_owned()),
275            max_message_size_bytes: None,
276            max_frame_size_bytes: None,
277        };
278
279        let keyed_quotas = vec![(
280            COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION.to_string(),
281            *COINBASE_WS_SUBSCRIPTION_QUOTA,
282        )];
283
284        let client = WebSocketClient::builder()
285            .config(cfg)
286            .message_handler(message_handler)
287            .keyed_quotas(keyed_quotas)
288            .default_quota(*COINBASE_WS_CONNECTION_QUOTA)
289            .maybe_state_sink(self.socket_control.as_ref().map(SocketControl::sink))
290            .connect()
291            .await?;
292
293        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
294        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<NautilusWsMessage>();
295
296        *self.cmd_tx.write().await = cmd_tx.clone();
297        self.out_rx = Some(out_rx);
298        self.connection_mode.store(client.connection_mode_atomic());
299        let reconnect_handle = client.reconnect_handle();
300        log::debug!("Coinbase WebSocket connected: {}", self.url.origin());
301
302        if let Err(e) = cmd_tx.send(HandlerCommand::SetClient(client)) {
303            anyhow::bail!("Failed to send SetClient command: {e}");
304        }
305
306        if let Some(control) = &self.socket_control {
307            control.register(move || reconnect_handle.request_reconnect());
308        }
309
310        let instruments_vec: Vec<InstrumentAny> =
311            self.instruments.load().values().cloned().collect();
312
313        if !instruments_vec.is_empty()
314            && let Err(e) = cmd_tx.send(HandlerCommand::InitializeInstruments(instruments_vec))
315        {
316            log::error!("Failed to send InitializeInstruments: {e}");
317        }
318
319        // Restore bar type registrations from previous session
320        for (key, bar_type) in &self.bar_types {
321            if let Err(e) = cmd_tx.send(HandlerCommand::AddBarType {
322                key: key.clone(),
323                bar_type: *bar_type,
324            }) {
325                log::error!("Failed to restore bar type {key}: {e}");
326            }
327        }
328
329        if let Some(account_id) = self.account_id
330            && let Err(e) = cmd_tx.send(HandlerCommand::SetAccountId(account_id))
331        {
332            log::error!("Failed to restore account_id: {e}");
333        }
334
335        self.prime_default_subscriptions();
336
337        // Replay retained subscriptions from previous session
338        resubscribe_all(
339            &self.subscriptions,
340            &self.credential,
341            &cmd_tx,
342            Some(&out_tx),
343        );
344
345        let signal = Arc::clone(&self.signal);
346        let subscriptions = self.subscriptions.clone();
347        let credential = self.credential.clone();
348        let cmd_tx_reconnect = cmd_tx.clone();
349        let aliases_for_handler = Arc::clone(&self.subscription_aliases);
350
351        if let Err(e) = self.task_handle.spawn(async move {
352            let mut handler = FeedHandler::new(signal, cmd_rx, raw_rx, aliases_for_handler);
353
354            loop {
355                match handler.next().await {
356                    Some(NautilusWsMessage::Reconnected) => {
357                        subscriptions.reset_after_reconnect();
358                        resubscribe_all(
359                            &subscriptions,
360                            &credential,
361                            &cmd_tx_reconnect,
362                            Some(&out_tx),
363                        );
364
365                        if let Err(e) = out_tx.send(NautilusWsMessage::Reconnected) {
366                            log::debug!("Output channel closed: {e}");
367                            break;
368                        }
369                    }
370                    Some(msg) => {
371                        if let Err(e) = out_tx.send(msg) {
372                            log::debug!("Output channel closed: {e}");
373                            break;
374                        }
375                    }
376                    None => {
377                        log::debug!("Feed handler stopped");
378                        break;
379                    }
380                }
381            }
382        }) {
383            self.out_rx = None;
384            anyhow::bail!("Failed to start Coinbase WebSocket handler task: {e}");
385        }
386
387        Ok(())
388    }
389
390    /// Subscribes to a channel for the given product IDs.
391    pub async fn subscribe(
392        &self,
393        channel: CoinbaseWsChannel,
394        product_ids: &[Ustr],
395    ) -> anyhow::Result<()> {
396        let jwt = if channel.requires_auth() {
397            let credential = self
398                .credential
399                .as_ref()
400                .ok_or_else(|| anyhow::anyhow!("Credentials required for {channel}"))?;
401            Some(credential.build_ws_jwt()?)
402        } else {
403            self.credential.as_ref().and_then(|c| c.build_ws_jwt().ok())
404        };
405
406        let sub = protect_subscription(CoinbaseWsSubscription {
407            msg_type: CoinbaseWsAction::Subscribe,
408            product_ids: product_ids.to_vec(),
409            channel,
410            jwt,
411        })?;
412
413        let channel_str = channel.as_ref();
414
415        if product_ids.is_empty() {
416            self.subscriptions.mark_subscribe(channel_str);
417        } else {
418            for product_id in product_ids {
419                let topic = format!("{channel_str}|{product_id}");
420                self.subscriptions.mark_subscribe(&topic);
421            }
422        }
423
424        let cmd_tx = self.cmd_tx.read().await;
425        cmd_tx
426            .send(HandlerCommand::Subscribe {
427                channel: sub.channel,
428                product_ids: sub.product_ids,
429                payload: sub.payload,
430            })
431            .map_err(|e| anyhow::anyhow!("Failed to send Subscribe command: {e}"))
432    }
433
434    /// Unsubscribes from a channel for the given product IDs.
435    pub async fn unsubscribe(
436        &self,
437        channel: CoinbaseWsChannel,
438        product_ids: &[Ustr],
439    ) -> anyhow::Result<()> {
440        let jwt = self.credential.as_ref().and_then(|c| c.build_ws_jwt().ok());
441
442        let unsub = protect_subscription(CoinbaseWsSubscription {
443            msg_type: CoinbaseWsAction::Unsubscribe,
444            product_ids: product_ids.to_vec(),
445            channel,
446            jwt,
447        })?;
448
449        let channel_str = channel.as_ref();
450
451        if product_ids.is_empty() {
452            self.subscriptions.mark_unsubscribe(channel_str);
453        } else {
454            for product_id in product_ids {
455                let topic = format!("{channel_str}|{product_id}");
456                self.subscriptions.mark_unsubscribe(&topic);
457            }
458        }
459
460        let cmd_tx = self.cmd_tx.read().await;
461        cmd_tx
462            .send(HandlerCommand::Unsubscribe {
463                channel: unsub.channel,
464                product_ids: unsub.product_ids,
465                payload: unsub.payload,
466            })
467            .map_err(|e| anyhow::anyhow!("Failed to send Unsubscribe command: {e}"))
468    }
469
470    /// Returns the next parsed message from the feed handler.
471    pub async fn next_message(&mut self) -> Option<NautilusWsMessage> {
472        self.out_rx.as_mut()?.recv().await
473    }
474
475    /// Disconnects the WebSocket and stops the feed handler.
476    pub(crate) fn begin_shutdown(&self) {
477        self.signal.store(true, Ordering::Release);
478    }
479
480    /// Disconnects the WebSocket and stops the feed handler.
481    pub async fn disconnect(&mut self) -> anyhow::Result<()> {
482        // Send Disconnect command before setting the signal so the handler
483        // processes it and calls notify_closed() on the inner WebSocket client
484        let cmd_tx = self.cmd_tx.read().await;
485
486        if let Err(e) = cmd_tx.send(HandlerCommand::Disconnect) {
487            log::debug!("Failed to send Disconnect command: {e}");
488        }
489        drop(cmd_tx);
490
491        // Release pairs with the handler's Acquire load; fallback for when
492        // the command channel is full or closed.
493        self.begin_shutdown();
494
495        if let Some(outcome) = finish_task(
496            &mut self.task_handle,
497            WS_DISCONNECT_TIMEOUT,
498            WS_DISCONNECT_TIMEOUT,
499        )
500        .await
501        {
502            match outcome {
503                TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted => {}
504                TaskJoinOutcome::Failed(error) => {
505                    self.shutdown_errors
506                        .push(format!("Coinbase WebSocket handler failed: {error}"));
507                }
508                TaskJoinOutcome::Incomplete => {
509                    self.shutdown_errors
510                        .push("Coinbase WebSocket handler did not stop after abort".to_string());
511                }
512            }
513        }
514
515        // Wait for the inner WebSocket's connection_mode atomic to reach Closed
516        // before returning. Without this, a subsequent connect() can observe a
517        // stale Active/Reconnect state and early-return, leaving out_rx unset
518        // and causing "WebSocket output receiver not available" on take.
519        let deadline = tokio::time::Instant::now() + Duration::from_secs(5);
520
521        loop {
522            let mode_ptr = self.connection_mode.load();
523
524            if ConnectionMode::from_u8(mode_ptr.load(Ordering::Relaxed)).is_closed() {
525                break;
526            }
527
528            if tokio::time::Instant::now() >= deadline {
529                self.shutdown_errors
530                    .push("Timed out waiting for WebSocket to reach Closed state".to_string());
531                break;
532            }
533
534            tokio::time::sleep(Duration::from_millis(20)).await;
535        }
536
537        if let Some(control) = &self.socket_control {
538            control.deregister();
539        }
540
541        if self.shutdown_errors.is_empty() {
542            Ok(())
543        } else {
544            let errors = std::mem::take(&mut self.shutdown_errors);
545            anyhow::bail!(errors.join("; "))
546        }
547    }
548
549    /// Returns true if the WebSocket connection is active.
550    #[must_use]
551    pub fn is_active(&self) -> bool {
552        let mode_ptr = self.connection_mode.load();
553        let mode_val = mode_ptr.load(Ordering::Relaxed);
554        ConnectionMode::from_u8(mode_val).is_active()
555    }
556
557    /// Returns true if the WebSocket is reconnecting after a transport drop.
558    #[must_use]
559    pub fn is_reconnecting(&self) -> bool {
560        let mode_ptr = self.connection_mode.load();
561        let mode_val = mode_ptr.load(Ordering::Relaxed);
562        ConnectionMode::from_u8(mode_val).is_reconnect()
563    }
564
565    /// Returns a reference to the instrument cache.
566    #[must_use]
567    pub fn instruments(&self) -> &Arc<AtomicMap<InstrumentId, InstrumentAny>> {
568        &self.instruments
569    }
570
571    /// Returns a reference to the canonical-to-subscribed alias map.
572    #[must_use]
573    pub fn subscription_aliases(&self) -> &Arc<AtomicMap<Ustr, Ustr>> {
574        &self.subscription_aliases
575    }
576
577    /// Records that inbound messages carrying `canonical` should be re-keyed to
578    /// `subscribed`. Caller is the data/exec client at subscribe or submit time
579    /// when the local product id differs from Coinbase's canonical alias.
580    pub fn register_subscription_alias(&self, canonical: Ustr, subscribed: Ustr) {
581        self.subscription_aliases.insert(canonical, subscribed);
582    }
583
584    /// Removes an alias registration. Safe to call if no entry exists.
585    pub fn unregister_subscription_alias(&self, canonical: &Ustr) {
586        self.subscription_aliases.remove(canonical);
587    }
588
589    /// Returns the subscription state.
590    #[must_use]
591    pub fn subscriptions(&self) -> &SubscriptionState {
592        &self.subscriptions
593    }
594
595    /// Updates an instrument in the cache and notifies the handler.
596    pub async fn update_instrument(&self, instrument: InstrumentAny) {
597        let id = instrument.id();
598        self.instruments.insert(id, instrument.clone());
599
600        let cmd_tx = self.cmd_tx.read().await;
601
602        if let Err(e) = cmd_tx.send(HandlerCommand::UpdateInstrument(Box::new(instrument))) {
603            log::debug!("Failed to send UpdateInstrument: {e}");
604        }
605    }
606
607    /// Takes the output message receiver, leaving `None` in its place.
608    ///
609    /// Used by the data client to move the receiver into a background consumption task.
610    pub fn take_out_rx(
611        &mut self,
612    ) -> Option<tokio::sync::mpsc::UnboundedReceiver<NautilusWsMessage>> {
613        self.out_rx.take()
614    }
615
616    /// Registers a bar type locally without notifying the handler.
617    ///
618    /// Used by the data client to persist registrations on the original client
619    /// before cloning for async command dispatch.
620    pub fn register_bar_type(&mut self, key: String, bar_type: BarType) {
621        self.bar_types.insert(key, bar_type);
622    }
623
624    /// Registers a bar type for candle parsing.
625    pub async fn add_bar_type(&mut self, key: String, bar_type: BarType) {
626        self.bar_types.insert(key.clone(), bar_type);
627
628        let cmd_tx = self.cmd_tx.read().await;
629
630        if let Err(e) = cmd_tx.send(HandlerCommand::AddBarType { key, bar_type }) {
631            log::debug!("Failed to send AddBarType: {e}");
632        }
633    }
634}
635
636impl Drop for CoinbaseWebSocketClient {
637    fn drop(&mut self) {
638        if let Some(handle) = self.task_handle.as_ref() {
639            self.signal.store(true, Ordering::Release);
640            handle.abort();
641        }
642    }
643}
644
645struct ProtectedSubscription {
646    channel: CoinbaseWsChannel,
647    product_ids: Vec<Ustr>,
648    payload: SecretString,
649}
650
651fn protect_subscription(
652    mut subscription: CoinbaseWsSubscription,
653) -> Result<ProtectedSubscription, serde_json::Error> {
654    let payload = SecretString::from(serde_json::to_string(&subscription)?);
655
656    Ok(ProtectedSubscription {
657        channel: subscription.channel,
658        product_ids: std::mem::take(&mut subscription.product_ids),
659        payload,
660    })
661}
662
663fn resubscribe_all(
664    subscriptions: &SubscriptionState,
665    credential: &Option<CoinbaseCredential>,
666    cmd_tx: &tokio::sync::mpsc::UnboundedSender<HandlerCommand>,
667    out_tx: Option<&tokio::sync::mpsc::UnboundedSender<NautilusWsMessage>>,
668) {
669    let topics = subscriptions.all_topics();
670
671    if topics.is_empty() {
672        log::debug!("No active subscriptions to restore");
673        return;
674    }
675
676    log::info!(
677        "Resubscribing to {} topics after reconnection",
678        topics.len()
679    );
680
681    for topic in topics {
682        let (channel, product_id) = match topic.split_once('|') {
683            Some((ch, pid)) => (ch, Some(pid)),
684            None => (topic.as_str(), None),
685        };
686
687        let channel_enum = match CoinbaseWsChannel::from_str(channel) {
688            Ok(ch) => ch,
689            Err(_) => {
690                log::warn!("Unknown channel in topic: {topic}");
691                continue;
692            }
693        };
694
695        let jwt = match credential.as_ref() {
696            Some(c) => match c.build_ws_jwt() {
697                Ok(token) => Some(token),
698                Err(e) => {
699                    if channel_enum.requires_auth() {
700                        let msg = format!(
701                            "JWT required for {channel} but build failed: {e}; topic {topic} not restored"
702                        );
703                        log::error!("{msg}");
704                        if let Some(tx) = out_tx {
705                            let _ = tx.send(NautilusWsMessage::Error(msg));
706                        }
707                        continue;
708                    }
709                    None
710                }
711            },
712            None => {
713                if channel_enum.requires_auth() {
714                    let msg = format!(
715                        "JWT required for {channel} but no credentials configured; topic {topic} not restored"
716                    );
717                    log::error!("{msg}");
718                    if let Some(tx) = out_tx {
719                        let _ = tx.send(NautilusWsMessage::Error(msg));
720                    }
721                    continue;
722                }
723                None
724            }
725        };
726
727        let product_ids = match product_id {
728            Some(pid) => vec![Ustr::from(pid)],
729            None => vec![],
730        };
731
732        let sub = CoinbaseWsSubscription {
733            msg_type: CoinbaseWsAction::Subscribe,
734            product_ids,
735            channel: channel_enum,
736            jwt,
737        };
738
739        let sub = match protect_subscription(sub) {
740            Ok(sub) => sub,
741            Err(e) => {
742                log::error!("Failed to serialize subscription for {topic}: {e}");
743                continue;
744            }
745        };
746
747        if let Err(e) = cmd_tx.send(HandlerCommand::Subscribe {
748            channel: sub.channel,
749            product_ids: sub.product_ids,
750            payload: sub.payload,
751        }) {
752            log::error!("Failed to resubscribe {topic}: {e}");
753        }
754    }
755}
756
757#[cfg(test)]
758mod tests {
759    use nautilus_network::websocket::SubscriptionState;
760    use rstest::rstest;
761
762    use super::*;
763
764    #[rstest]
765    #[case("wss://stream.example.com/feed", "wss://stream.example.com")]
766    #[case(
767        "wss://SYN-user:SYN-password@stream.example.com:9443/SYN-path?token=SYN-query#SYN-fragment",
768        "wss://stream.example.com:9443"
769    )]
770    fn test_debug_preserves_url_origin(#[case] url: &str, #[case] origin: &str) {
771        let client = CoinbaseWebSocketClient::new(url, TransportBackend::default(), None);
772        let debug = format!("{client:?}");
773        assert!(debug.contains(&format!("url: {origin:?}")));
774        assert!(!debug.contains("SYN-"));
775        assert_eq!(client.url.expose(), url);
776    }
777
778    #[rstest]
779    fn test_debug_redacts_proxy_url() {
780        let proxy_url = "http://user:password@proxy.example:8080";
781        let client = CoinbaseWebSocketClient::new(
782            "wss://test",
783            TransportBackend::default(),
784            Some(proxy_url.to_string()),
785        );
786
787        let debug = format!("{client:?}");
788
789        assert!(debug.contains(REDACTED));
790        assert!(!debug.contains(proxy_url));
791    }
792
793    #[rstest]
794    fn test_drop_clone_does_not_signal_handler() {
795        let client = CoinbaseWebSocketClient::new("wss://test", TransportBackend::default(), None);
796        let clone = client.clone();
797
798        drop(clone);
799
800        assert!(!client.signal.load(Ordering::Acquire));
801    }
802
803    #[rstest]
804    fn test_resubscribe_all_product_level_topic() {
805        let subs = SubscriptionState::new('|');
806        subs.mark_subscribe("level2|BTC-USD");
807
808        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
809        resubscribe_all(&subs, &None, &tx, None);
810
811        let cmd = rx.try_recv().unwrap();
812
813        match cmd {
814            HandlerCommand::Subscribe {
815                channel,
816                product_ids,
817                payload,
818            } => {
819                let payload: serde_json::Value =
820                    serde_json::from_str(payload.expose_secret()).unwrap();
821                assert_eq!(channel, CoinbaseWsChannel::Level2);
822                assert_eq!(product_ids, vec![Ustr::from("BTC-USD")]);
823                assert!(payload.get("jwt").is_none());
824            }
825            other => panic!("Expected Subscribe, was {other:?}"),
826        }
827    }
828
829    #[rstest]
830    fn test_resubscribe_all_channel_level_topic() {
831        let subs = SubscriptionState::new('|');
832        subs.mark_subscribe("heartbeats");
833
834        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
835        resubscribe_all(&subs, &None, &tx, None);
836
837        let cmd = rx.try_recv().unwrap();
838
839        match cmd {
840            HandlerCommand::Subscribe {
841                channel,
842                product_ids,
843                ..
844            } => {
845                assert_eq!(channel, CoinbaseWsChannel::Heartbeats);
846                assert!(product_ids.is_empty());
847            }
848            other => panic!("Expected Subscribe, was {other:?}"),
849        }
850    }
851
852    #[rstest]
853    fn test_resubscribe_all_multiple_topics() {
854        let subs = SubscriptionState::new('|');
855        subs.mark_subscribe("market_trades|BTC-USD");
856        subs.mark_subscribe("ticker|ETH-USD");
857
858        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
859        resubscribe_all(&subs, &None, &tx, None);
860
861        let cmd1 = rx.try_recv().unwrap();
862        let cmd2 = rx.try_recv().unwrap();
863
864        assert!(matches!(cmd1, HandlerCommand::Subscribe { .. }));
865        assert!(matches!(cmd2, HandlerCommand::Subscribe { .. }));
866        assert!(rx.try_recv().is_err());
867    }
868
869    #[rstest]
870    fn test_resubscribe_all_empty_subscriptions() {
871        let subs = SubscriptionState::new('|');
872
873        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
874        resubscribe_all(&subs, &None, &tx, None);
875
876        assert!(rx.try_recv().is_err());
877    }
878
879    #[rstest]
880    fn test_resubscribe_all_unknown_channel_skipped() {
881        let subs = SubscriptionState::new('|');
882        subs.mark_subscribe("nonexistent_channel|BTC-USD");
883
884        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
885        resubscribe_all(&subs, &None, &tx, None);
886
887        assert!(rx.try_recv().is_err());
888    }
889
890    #[rstest]
891    #[case("level2|BTC-USD", CoinbaseWsChannel::Level2)]
892    #[case("market_trades|ETH-USD", CoinbaseWsChannel::MarketTrades)]
893    #[case("ticker|BTC-USD", CoinbaseWsChannel::Ticker)]
894    #[case("ticker_batch|BTC-USD", CoinbaseWsChannel::TickerBatch)]
895    #[case("candles|BTC-USD", CoinbaseWsChannel::Candles)]
896    #[case("heartbeats", CoinbaseWsChannel::Heartbeats)]
897    #[case("status", CoinbaseWsChannel::Status)]
898    fn test_resubscribe_all_channel_mapping(
899        #[case] topic: &str,
900        #[case] expected_channel: CoinbaseWsChannel,
901    ) {
902        let subs = SubscriptionState::new('|');
903        subs.mark_subscribe(topic);
904
905        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
906        resubscribe_all(&subs, &None, &tx, None);
907
908        let cmd = rx.try_recv().unwrap();
909
910        match cmd {
911            HandlerCommand::Subscribe { channel, .. } => {
912                assert_eq!(channel, expected_channel);
913            }
914            other => panic!("Expected Subscribe, was {other:?}"),
915        }
916    }
917
918    #[rstest]
919    fn test_protected_subscription_preserves_wire_jwt_and_redacts_command_debug() {
920        let jwt = "jwt-secret-value";
921        let subscription = CoinbaseWsSubscription {
922            msg_type: CoinbaseWsAction::Subscribe,
923            product_ids: vec![Ustr::from("BTC-USD")],
924            channel: CoinbaseWsChannel::User,
925            jwt: Some(SecretString::from(jwt)),
926        };
927        let subscription_debug = format!("{subscription:?}");
928        let protected = protect_subscription(subscription).unwrap();
929        let payload: serde_json::Value =
930            serde_json::from_str(protected.payload.expose_secret()).unwrap();
931        let command = HandlerCommand::Subscribe {
932            channel: protected.channel,
933            product_ids: protected.product_ids,
934            payload: protected.payload,
935        };
936        let debug = format!("{command:?}");
937
938        assert_eq!(payload["jwt"], jwt);
939        assert!(subscription_debug.contains(REDACTED));
940        assert!(!subscription_debug.contains(jwt));
941        assert!(!debug.contains(jwt));
942    }
943
944    #[rstest]
945    #[case("user|BTC-USD")]
946    #[case("futures_balance_summary")]
947    fn test_resubscribe_all_auth_channel_skipped_without_credentials(#[case] topic: &str) {
948        let subs = SubscriptionState::new('|');
949        subs.mark_subscribe(topic);
950
951        let (tx, mut rx) = tokio::sync::mpsc::unbounded_channel();
952        resubscribe_all(&subs, &None, &tx, None);
953
954        // Auth channels should be skipped when no credentials are provided
955        assert!(rx.try_recv().is_err());
956    }
957
958    #[rstest]
959    #[case("user|BTC-USD", "user")]
960    #[case("futures_balance_summary", "futures_balance_summary")]
961    fn test_resubscribe_all_emits_error_for_auth_channel_without_credentials(
962        #[case] topic: &str,
963        #[case] channel: &str,
964    ) {
965        let subs = SubscriptionState::new('|');
966        subs.mark_subscribe(topic);
967
968        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
969        let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
970        resubscribe_all(&subs, &None, &cmd_tx, Some(&out_tx));
971
972        // No subscribe command should be sent for an unauthenticated auth channel.
973        assert!(cmd_rx.try_recv().is_err());
974
975        let msg = out_rx
976            .try_recv()
977            .expect("Error event must be emitted when auth channel cannot resubscribe");
978        match msg {
979            NautilusWsMessage::Error(text) => {
980                assert!(
981                    text.contains(channel),
982                    "error must mention the channel, was: {text}"
983                );
984                assert!(
985                    text.contains(topic),
986                    "error must mention the topic, was: {text}"
987                );
988            }
989            other => panic!("expected Error variant, was {other:?}"),
990        }
991    }
992
993    #[rstest]
994    fn test_resubscribe_all_emits_error_when_jwt_build_fails() {
995        let subs = SubscriptionState::new('|');
996        let topic = "user|BTC-USD";
997        subs.mark_subscribe(topic);
998
999        // A credential with a malformed PEM secret causes build_ws_jwt() to fail
1000        // every time, exercising the JWT-build error branch.
1001        let bad_credential = Some(CoinbaseCredential::new(
1002            "organizations/test/apiKeys/test".to_string(),
1003            "not-a-pem-key".to_string(),
1004        ));
1005
1006        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
1007        let (out_tx, mut out_rx) = tokio::sync::mpsc::unbounded_channel();
1008        resubscribe_all(&subs, &bad_credential, &cmd_tx, Some(&out_tx));
1009
1010        assert!(cmd_rx.try_recv().is_err(), "no subscribe should be sent");
1011        let msg = out_rx
1012            .try_recv()
1013            .expect("Error event must be emitted when JWT build fails for an auth channel");
1014        match msg {
1015            NautilusWsMessage::Error(text) => {
1016                assert!(text.contains("user"), "error must mention channel: {text}");
1017                assert!(text.contains(topic), "error must mention topic: {text}");
1018            }
1019            other => panic!("expected Error variant, was {other:?}"),
1020        }
1021    }
1022
1023    #[rstest]
1024    fn test_prime_default_subscriptions_marks_heartbeats() {
1025        let client = CoinbaseWebSocketClient::new("wss://test", TransportBackend::default(), None);
1026        assert!(client.subscriptions.all_topics().is_empty());
1027
1028        client.prime_default_subscriptions();
1029
1030        let topics = client.subscriptions.all_topics();
1031        assert!(topics.iter().any(|t| t == "heartbeats"), "{topics:?}");
1032    }
1033
1034    #[rstest]
1035    fn test_ws_quotas_match_documented_limits() {
1036        assert_eq!(COINBASE_WS_CONNECTION_QUOTA.burst_size().get(), 8);
1037        assert_eq!(COINBASE_WS_SUBSCRIPTION_QUOTA.burst_size().get(), 8);
1038    }
1039
1040    #[rstest]
1041    fn test_ws_subscription_rate_limit_key_is_stable() {
1042        assert_eq!(COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION, "subscription");
1043        assert_eq!(
1044            COINBASE_WS_SUBSCRIPTION_KEYS[0],
1045            COINBASE_RATE_LIMIT_KEY_SUBSCRIPTION,
1046        );
1047    }
1048}