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