nautilus_polymarket/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//! Provides the WebSocket client for the Polymarket CLOB API.
17
18use std::sync::{
19    Arc,
20    atomic::{AtomicBool, AtomicU8, Ordering},
21};
22
23use nautilus_live::{
24    SocketControl,
25    task::{TaskJoinOutcome, TaskSlot, finish_task},
26};
27use nautilus_network::{
28    SocketStateSink,
29    mode::ConnectionMode,
30    websocket::{
31        AuthTracker, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
32        channel_epoch_message_handler, proxy::ProxyUrl,
33    },
34};
35
36use super::{
37    handler::{FeedHandler, HandlerCommand},
38    messages::PolymarketWsMessage,
39};
40use crate::common::{
41    credential::Credential,
42    urls::{clob_ws_market_url, clob_ws_user_url},
43};
44
45// The venue counts only the `PING` text frame, not protocol ping frames, and
46// closes with `1008 no ping received` otherwise. Cadence per venue docs:
47// https://docs.polymarket.com/api-reference/wss/market
48pub(super) const POLYMARKET_HEARTBEAT_SECS: u64 = 10;
49pub(super) const POLYMARKET_HEARTBEAT_PAYLOAD: &str = "PING";
50
51// Prediction markets go quiet for long stretches, so liveness is the venue
52// still sending frames, not data arriving. A data-silence timer cannot serve:
53// `PONG` is a text frame and refreshes it. Tear down after three cycles.
54const POLYMARKET_HEARTBEAT_TIMEOUT_SECS: u64 = POLYMARKET_HEARTBEAT_SECS * 3;
55
56/// Polymarket WebSocket channel: market data or authenticated user data.
57#[derive(Debug, Clone, Copy, PartialEq, Eq)]
58pub enum WsChannel {
59    Market,
60    User,
61}
62
63/// Lightweight handle for subscribing/unsubscribing to market data.
64///
65/// `Clone` + `Send` safe for use in spawned async tasks.
66#[derive(Clone, Debug)]
67pub struct WsSubscriptionHandle {
68    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
69}
70
71impl WsSubscriptionHandle {
72    /// Sends a market subscribe command to the handler.
73    pub async fn subscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
74        self.cmd_tx
75            .read()
76            .await
77            .send(HandlerCommand::SubscribeMarket(asset_ids))
78            .map_err(|e| anyhow::anyhow!("Failed to send SubscribeMarket: {e}"))
79    }
80
81    /// Sends a market unsubscribe command to the handler.
82    pub async fn unsubscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
83        self.cmd_tx
84            .read()
85            .await
86            .send(HandlerCommand::UnsubscribeMarket(asset_ids))
87            .map_err(|e| anyhow::anyhow!("Failed to send UnsubscribeMarket: {e}"))
88    }
89
90    // Constructs a handle around a raw command sender. Test-only: lets unit
91    // tests observe the commands the handle emits without spinning up the real
92    // feed handler.
93    #[cfg(test)]
94    pub(crate) fn from_sender(sender: tokio::sync::mpsc::UnboundedSender<HandlerCommand>) -> Self {
95        Self {
96            cmd_tx: Arc::new(tokio::sync::RwLock::new(sender)),
97        }
98    }
99}
100
101/// Provides a WebSocket client for the Polymarket CLOB API.
102///
103/// A single instance targets one channel (market or user). Use
104/// [`PolymarketWebSocketClient::new_market`] for public market data and
105/// [`PolymarketWebSocketClient::new_user`] for authenticated order/trade streams.
106#[derive(Debug)]
107pub struct PolymarketWebSocketClient {
108    channel: WsChannel,
109    url: String,
110    connection_mode: Arc<AtomicU8>,
111    signal: Arc<AtomicBool>,
112    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
113    out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>>,
114    credential: Option<Credential>,
115    subscriptions: SubscriptionState,
116    discovery_subscribed: Arc<AtomicBool>,
117    auth_tracker: AuthTracker,
118    // Survives disconnect() so that connect() can replay a prior subscribe_user() call.
119    // Arc<AtomicBool> allows mutation from &self in subscribe_user().
120    user_subscribed: Arc<AtomicBool>,
121    task_handle: TaskSlot<()>,
122    subscribe_new_markets: bool,
123    transport_backend: TransportBackend,
124    proxy_url: Option<ProxyUrl>,
125    socket_sink: Option<SocketStateSink>,
126    socket_control: Option<SocketControl>,
127}
128
129#[derive(Clone, Debug)]
130pub(crate) struct PolymarketWebSocketShutdownHandle {
131    signal: Arc<AtomicBool>,
132}
133
134impl PolymarketWebSocketShutdownHandle {
135    pub(crate) fn begin_shutdown(&self) {
136        self.signal.store(true, Ordering::Relaxed);
137    }
138}
139
140impl PolymarketWebSocketClient {
141    /// Creates a new market-channel client (unauthenticated).
142    ///
143    /// If `base_url` is `None`, the default production URL is used.
144    #[must_use]
145    pub fn new_market(
146        base_url: Option<String>,
147        subscribe_new_markets: bool,
148        transport_backend: TransportBackend,
149    ) -> Self {
150        Self::new_market_with_proxy(base_url, subscribe_new_markets, transport_backend, None)
151    }
152
153    /// Creates a new market-channel client with an optional validated proxy URL.
154    #[must_use]
155    pub fn new_market_with_proxy(
156        base_url: Option<String>,
157        subscribe_new_markets: bool,
158        transport_backend: TransportBackend,
159        proxy_url: Option<ProxyUrl>,
160    ) -> Self {
161        let url = base_url.unwrap_or_else(|| clob_ws_market_url().to_string());
162        Self::new_inner(
163            WsChannel::Market,
164            url,
165            None,
166            subscribe_new_markets,
167            transport_backend,
168            proxy_url,
169        )
170    }
171
172    /// Creates a new user-channel client (authenticated).
173    ///
174    /// If `base_url` is `None`, the default production URL is used.
175    #[must_use]
176    pub fn new_user(
177        base_url: Option<String>,
178        credential: Credential,
179        transport_backend: TransportBackend,
180    ) -> Self {
181        Self::new_user_with_proxy(base_url, credential, transport_backend, None)
182    }
183
184    /// Creates a new user-channel client with an optional validated proxy URL.
185    #[must_use]
186    pub fn new_user_with_proxy(
187        base_url: Option<String>,
188        credential: Credential,
189        transport_backend: TransportBackend,
190        proxy_url: Option<ProxyUrl>,
191    ) -> Self {
192        let url = base_url.unwrap_or_else(|| clob_ws_user_url().to_string());
193        Self::new_inner(
194            WsChannel::User,
195            url,
196            Some(credential),
197            false,
198            transport_backend,
199            proxy_url,
200        )
201    }
202
203    fn new_inner(
204        channel: WsChannel,
205        url: String,
206        credential: Option<Credential>,
207        subscribe_new_markets: bool,
208        transport_backend: TransportBackend,
209        proxy_url: Option<ProxyUrl>,
210    ) -> Self {
211        let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
212        Self {
213            channel,
214            url,
215            connection_mode: Arc::new(AtomicU8::new(ConnectionMode::Closed.as_u8())),
216            signal: Arc::new(AtomicBool::new(false)),
217            cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
218            out_rx: None,
219            credential,
220            subscriptions: SubscriptionState::new(':'),
221            discovery_subscribed: Arc::new(AtomicBool::new(false)),
222            auth_tracker: AuthTracker::new(),
223            user_subscribed: Arc::new(AtomicBool::new(false)),
224            task_handle: TaskSlot::new(),
225            subscribe_new_markets,
226            transport_backend,
227            proxy_url,
228            socket_sink: None,
229            socket_control: None,
230        }
231    }
232
233    /// Configures socket state reporting for the underlying transport.
234    #[must_use]
235    pub fn with_state_sink(mut self, state_sink: SocketStateSink) -> Self {
236        self.socket_sink = Some(state_sink);
237        self
238    }
239
240    /// Configures state reporting and reconnect control for the underlying transport.
241    #[must_use]
242    pub(crate) fn with_socket_control(mut self, control: SocketControl) -> Self {
243        self.socket_control = Some(control);
244        self
245    }
246
247    #[cfg(test)]
248    pub(crate) fn proxy_url(&self) -> Option<&ProxyUrl> {
249        self.proxy_url.as_ref()
250    }
251
252    /// Establishes the WebSocket connection and spawns the message handler.
253    pub async fn connect(&mut self) -> anyhow::Result<()> {
254        let mode = ConnectionMode::from_atomic(&self.connection_mode);
255        if mode.is_active() || mode.is_reconnect() {
256            log::warn!("Polymarket WebSocket already connected or reconnecting");
257            return Ok(());
258        }
259
260        if self.task_handle.is_some() {
261            self.disconnect().await?;
262        }
263
264        let (message_handler, raw_rx) = channel_epoch_message_handler();
265        let cfg = self.websocket_config();
266
267        let client = WebSocketClient::epoch_builder()
268            .config(cfg)
269            .epoch_handler(message_handler)
270            .maybe_state_sink(
271                self.socket_control
272                    .as_ref()
273                    .map(SocketControl::sink)
274                    .or_else(|| self.socket_sink.clone()),
275            )
276            .connect()
277            .await?;
278
279        if let Some(control) = &self.socket_control {
280            let handle = client.reconnect_handle();
281            control.register(move || handle.request_reconnect());
282        }
283        let connection_epoch = client.connection_epoch();
284
285        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
286        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<PolymarketWsMessage>();
287
288        *self.cmd_tx.write().await = cmd_tx.clone();
289        self.out_rx = Some(out_rx);
290
291        let client_mode = client.connection_mode_atomic();
292        self.connection_mode = client_mode;
293
294        log::debug!("Polymarket WebSocket connected: {}", self.url);
295
296        // Replay retained state onto the new session. Unlike the RECONNECTED sentinel
297        // path, a fresh connect() never fires resubscribe_all() inside the handler.
298        let initial_market_replay = match self.channel {
299            WsChannel::Market => {
300                let topics = self.subscriptions.reset_after_reconnect();
301                if !topics.is_empty() || self.discovery_subscribed.load(Ordering::Relaxed) {
302                    log::debug!(
303                        "Replaying market subscription state onto new session: assets={}, discovery={}",
304                        topics.len(),
305                        self.discovery_subscribed.load(Ordering::Relaxed),
306                    );
307                    Some((topics, connection_epoch))
308                } else {
309                    None
310                }
311            }
312            WsChannel::User => {
313                if self.user_subscribed.load(Ordering::Relaxed) {
314                    log::debug!("Replaying user subscribe onto new session");
315                    cmd_tx
316                        .send(HandlerCommand::SubscribeUser)
317                        .map_err(|e| anyhow::anyhow!("Failed to replay SubscribeUser: {e}"))?;
318                }
319                None
320            }
321        };
322
323        let signal = Arc::clone(&self.signal);
324        let channel = self.channel;
325        let credential = self.credential.clone();
326        let subscriptions = self.subscriptions.clone();
327        let discovery_subscribed = Arc::clone(&self.discovery_subscribed);
328        let auth_tracker = self.auth_tracker.clone();
329        let user_subscribed = self.user_subscribed.load(Ordering::Relaxed);
330        let subscribe_new_markets = self.subscribe_new_markets;
331
332        if let Err(e) = self.task_handle.spawn(async move {
333            let mut handler = FeedHandler::new(
334                signal,
335                channel,
336                Some(client),
337                cmd_rx,
338                raw_rx,
339                out_tx,
340                credential,
341                subscriptions,
342                discovery_subscribed,
343                initial_market_replay,
344                auth_tracker,
345                user_subscribed,
346                subscribe_new_markets,
347            );
348
349            loop {
350                match handler.next().await {
351                    Some(PolymarketWsMessage::Reconnected) => {
352                        log::info!("Polymarket WebSocket reconnected");
353
354                        if handler.send(PolymarketWsMessage::Reconnected).is_err() {
355                            if handler.is_stopped() {
356                                log::debug!("Output channel closed, stopping handler");
357                            } else {
358                                log::error!("Output channel closed, stopping handler");
359                            }
360                            break;
361                        }
362                    }
363                    Some(msg) => {
364                        if handler.send(msg).is_err() {
365                            if handler.is_stopped() {
366                                log::debug!("Output channel closed, stopping handler");
367                            } else {
368                                log::error!("Output channel closed, stopping handler");
369                            }
370                            break;
371                        }
372                    }
373                    None => {
374                        if handler.is_stopped() {
375                            log::debug!("Stop signal received, ending handler task");
376                        } else {
377                            log::warn!("Polymarket WebSocket stream ended unexpectedly");
378                        }
379                        break;
380                    }
381                }
382            }
383            log::debug!("Polymarket WebSocket handler task completed");
384        }) {
385            self.out_rx = None;
386            anyhow::bail!("Failed to start Polymarket WebSocket handler task: {e}");
387        }
388        Ok(())
389    }
390
391    fn websocket_config(&self) -> WebSocketConfig {
392        // The market endpoint rejects text PING before its initial subscription. Protocol pings
393        // keep an idle socket alive until FeedHandler starts the required text heartbeat.
394        let heartbeat_payload = match self.channel {
395            WsChannel::Market => None,
396            WsChannel::User => Some(POLYMARKET_HEARTBEAT_PAYLOAD.to_string()),
397        };
398
399        WebSocketConfig {
400            url: self.url.clone(),
401            headers: vec![],
402            heartbeat_interval_secs: Some(POLYMARKET_HEARTBEAT_SECS),
403            heartbeat_payload,
404            connect_timeout_ms: Some(15_000),
405            reconnect_delay_initial_ms: Some(250),
406            reconnect_delay_max_ms: Some(5_000),
407            reconnect_backoff_factor: Some(2.0),
408            reconnect_jitter_ms: Some(200),
409            reconnect_max_attempts: None,
410            heartbeat_timeout_secs: Some(POLYMARKET_HEARTBEAT_TIMEOUT_SECS),
411            idle_timeout_ms: None,
412            backend: self.transport_backend,
413            proxy_url: self.proxy_url.as_ref().map(|url| url.expose().to_string()),
414        }
415    }
416
417    pub(crate) fn begin_shutdown(&self) {
418        self.signal.store(true, Ordering::Relaxed);
419    }
420
421    pub(crate) fn shutdown_handle(&self) -> PolymarketWebSocketShutdownHandle {
422        PolymarketWebSocketShutdownHandle {
423            signal: Arc::clone(&self.signal),
424        }
425    }
426
427    pub(crate) fn abort(&mut self) {
428        self.begin_shutdown();
429        self.connection_mode
430            .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
431        self.task_handle.abort();
432        self.auth_tracker.invalidate();
433
434        if let Some(control) = &self.socket_control {
435            control.deregister();
436        }
437    }
438
439    /// Disconnects the WebSocket connection.
440    pub async fn disconnect(&mut self) -> anyhow::Result<()> {
441        log::debug!("Disconnecting Polymarket WebSocket");
442        self.signal.store(true, Ordering::Relaxed);
443
444        if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
445            log::debug!("Failed to send disconnect (handler may already be shut down): {e}");
446        }
447
448        let task_result = match finish_task(
449            &mut self.task_handle,
450            std::time::Duration::from_secs(2),
451            std::time::Duration::from_secs(2),
452        )
453        .await
454        {
455            None | Some(TaskJoinOutcome::Completed(()) | TaskJoinOutcome::Aborted) => Ok(()),
456            Some(TaskJoinOutcome::Failed(error)) => Err(anyhow::anyhow!(
457                "Polymarket WebSocket handler failed: {error}"
458            )),
459            Some(TaskJoinOutcome::Incomplete) => Err(anyhow::anyhow!(
460                "Polymarket WebSocket handler did not stop after abort"
461            )),
462        };
463        // Invalidate after the task has stopped so any in-flight auth_tracker.succeed()
464        // calls from the handler cannot race with and survive the invalidation.
465        self.auth_tracker.invalidate();
466
467        if let Some(control) = &self.socket_control {
468            control.deregister();
469        }
470        log::debug!("Polymarket WebSocket disconnected");
471        task_result
472    }
473
474    /// Returns `true` if the WebSocket is actively connected.
475    #[must_use]
476    pub fn is_active(&self) -> bool {
477        ConnectionMode::from_atomic(&self.connection_mode).is_active()
478    }
479
480    pub(crate) fn has_task(&self) -> bool {
481        self.task_handle.is_some()
482    }
483
484    /// Returns the URL this client connects to.
485    #[must_use]
486    pub fn url(&self) -> &str {
487        &self.url
488    }
489
490    /// Returns the number of active market asset subscriptions (pending + confirmed).
491    #[must_use]
492    pub fn subscription_count(&self) -> usize {
493        self.subscriptions.all_topics().len()
494    }
495
496    /// Clears retained subscription/auth replay state.
497    ///
498    /// Useful for hard resets where the caller wants reconnect to start from a
499    /// clean slate rather than replaying a previous generation's topics.
500    pub(crate) fn clear_reconnect_state(&self) {
501        self.subscriptions.clear();
502        self.discovery_subscribed.store(false, Ordering::Relaxed);
503        self.user_subscribed.store(false, Ordering::Relaxed);
504        self.auth_tracker.invalidate();
505    }
506
507    /// Returns `true` if the user channel has been authenticated.
508    #[must_use]
509    pub fn is_authenticated(&self) -> bool {
510        self.auth_tracker.is_authenticated()
511    }
512
513    /// Subscribe to market data for the given asset IDs.
514    ///
515    /// Sends a subscribe message immediately if connected; the IDs are also
516    /// retained so they are re-sent automatically on reconnect.
517    ///
518    /// # Errors
519    ///
520    /// Returns an error if called on a user-channel client (incompatible channel).
521    pub async fn subscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
522        if self.channel != WsChannel::Market {
523            anyhow::bail!(
524                "subscribe_market() requires a market-channel client (created with new_market())"
525            );
526        }
527        self.cmd_tx
528            .read()
529            .await
530            .send(HandlerCommand::SubscribeMarket(asset_ids))
531            .map_err(|e| anyhow::anyhow!("Failed to send SubscribeMarket: {e}"))
532    }
533
534    /// Remove asset IDs from the active subscription set.
535    ///
536    /// The IDs are dropped from the reconnect set so they will not be
537    /// re-subscribed after a reconnect. No wire message is sent.
538    ///
539    /// # Errors
540    ///
541    /// Returns an error if called on a user-channel client (incompatible channel).
542    pub async fn unsubscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
543        if self.channel != WsChannel::Market {
544            anyhow::bail!(
545                "unsubscribe_market() requires a market-channel client (created with new_market())"
546            );
547        }
548        self.cmd_tx
549            .read()
550            .await
551            .send(HandlerCommand::UnsubscribeMarket(asset_ids))
552            .map_err(|e| anyhow::anyhow!("Failed to send UnsubscribeMarket: {e}"))
553    }
554
555    /// Authenticate and subscribe to the user channel.
556    ///
557    /// # Errors
558    ///
559    /// Returns an error if called on a market-channel client (no credentials available).
560    pub async fn subscribe_user(&self) -> anyhow::Result<()> {
561        if self.channel != WsChannel::User {
562            anyhow::bail!(
563                "subscribe_user() requires a user-channel client (created with new_user())"
564            );
565        }
566        self.cmd_tx
567            .read()
568            .await
569            .send(HandlerCommand::SubscribeUser)
570            .map_err(|e| anyhow::anyhow!("Failed to send SubscribeUser: {e}"))?;
571        // Set only after the command is successfully enqueued so a failed send does not
572        // leave user_subscribed=true and cause an unintended replay on the next connect().
573        self.user_subscribed.store(true, Ordering::Relaxed);
574        Ok(())
575    }
576
577    /// Returns a cloneable subscription handle for use in spawned tasks.
578    #[must_use]
579    pub fn clone_subscription_handle(&self) -> WsSubscriptionHandle {
580        WsSubscriptionHandle {
581            cmd_tx: Arc::clone(&self.cmd_tx),
582        }
583    }
584
585    /// Takes the message receiver, leaving `None` in its place.
586    ///
587    /// This is useful when the data client needs to spawn its own handler
588    /// task that reads messages independently of the WS client.
589    /// Subscription methods (`subscribe_market`, etc.) remain usable on `&self`.
590    #[must_use]
591    pub fn take_message_receiver(
592        &mut self,
593    ) -> Option<tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>> {
594        self.out_rx.take()
595    }
596
597    /// Receives the next message from the WebSocket handler.
598    ///
599    /// Returns `None` when the handler has disconnected or the receiver
600    /// was not yet initialized (call `connect` first).
601    pub async fn next_message(&mut self) -> Option<PolymarketWsMessage> {
602        if let Some(ref mut rx) = self.out_rx {
603            rx.recv().await
604        } else {
605            None
606        }
607    }
608}
609
610impl Drop for PolymarketWebSocketClient {
611    fn drop(&mut self) {
612        self.signal.store(true, Ordering::Relaxed);
613
614        if let Some(handle) = self.task_handle.as_ref() {
615            handle.abort();
616        }
617
618        if let Some(control) = &self.socket_control {
619            control.deregister();
620        }
621    }
622}
623
624#[cfg(test)]
625mod tests {
626    use std::net::SocketAddr;
627
628    use axum::{
629        Router,
630        extract::ws::{Message as AxumWsMessage, WebSocket, WebSocketUpgrade},
631        response::Response,
632        routing::get,
633    };
634    use nautilus_network::{
635        RECONNECTED,
636        websocket::{TransportBackend, WebSocketConfig, proxy::ProxyUrl},
637    };
638    use rstest::rstest;
639
640    use super::*;
641
642    async fn handle_upgrade(ws: WebSocketUpgrade) -> Response {
643        ws.on_upgrade(handle_socket)
644    }
645
646    async fn handle_socket(mut socket: WebSocket) {
647        let _ = socket
648            .send(AxumWsMessage::Text(RECONNECTED.to_string().into()))
649            .await;
650    }
651
652    async fn start_test_server() -> SocketAddr {
653        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
654            .await
655            .expect("bind test websocket server");
656        let addr = listener.local_addr().expect("test websocket address");
657        let router = Router::new().route("/ws", get(handle_upgrade));
658
659        tokio::spawn(async move {
660            axum::serve(listener, router)
661                .await
662                .expect("test websocket server failed");
663        });
664
665        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
666        addr
667    }
668
669    #[tokio::test]
670    async fn cancelled_disconnect_retains_handler_task() {
671        let mut client = PolymarketWebSocketClient::new_market(
672            Some("ws://127.0.0.1:0".to_string()),
673            false,
674            TransportBackend::default(),
675        );
676        let (cmd_tx, mut cmd_rx) = tokio::sync::mpsc::unbounded_channel();
677        client.cmd_tx = Arc::new(tokio::sync::RwLock::new(cmd_tx));
678        client
679            .task_handle
680            .insert(tokio::spawn(std::future::pending()));
681
682        {
683            let disconnect = client.disconnect();
684            tokio::pin!(disconnect);
685            tokio::select! {
686                result = &mut disconnect => panic!("disconnect completed unexpectedly: {result:?}"),
687                command = cmd_rx.recv() => assert!(command.is_some()),
688            }
689        }
690
691        assert!(client.task_handle.is_some());
692    }
693
694    #[rstest]
695    #[tokio::test]
696    async fn connect_forwards_reconnected_message_to_receiver() {
697        let addr = start_test_server().await;
698        let mut client = PolymarketWebSocketClient::new_market(
699            Some(format!("ws://{addr}/ws")),
700            false,
701            TransportBackend::default(),
702        );
703
704        client.connect().await.expect("connect websocket client");
705
706        let message =
707            tokio::time::timeout(tokio::time::Duration::from_secs(2), client.next_message())
708                .await
709                .expect("wait for websocket message");
710
711        assert!(matches!(
712            message,
713            Some(super::super::messages::PolymarketWsMessage::Reconnected)
714        ));
715
716        client
717            .disconnect()
718            .await
719            .expect("disconnect websocket client");
720    }
721
722    #[rstest]
723    fn proxy_url_is_retained_for_market_and_user_clients() {
724        const MARKET_PROXY: &str = "http://market-user:market-proxy-secret@127.0.0.1:18086";
725        const USER_PROXY: &str = "https://user-user:user-proxy-secret@127.0.0.1:18087";
726        let market = PolymarketWebSocketClient::new_market_with_proxy(
727            Some("ws://market.example/ws".to_string()),
728            false,
729            TransportBackend::Tungstenite,
730            Some(ProxyUrl::parse(MARKET_PROXY).unwrap()),
731        );
732        let credential = crate::common::credential::Credential::new(
733            "fixture-key",
734            "Zml4dHVyZQ==",
735            "fixture-passphrase".to_string(),
736        )
737        .unwrap();
738        let user = PolymarketWebSocketClient::new_user_with_proxy(
739            Some("ws://user.example/ws".to_string()),
740            credential,
741            TransportBackend::Tungstenite,
742            Some(ProxyUrl::parse(USER_PROXY).unwrap()),
743        );
744        let market_config = market.websocket_config();
745        let user_config = user.websocket_config();
746        let market_debug = format!("{market:?}");
747        let user_debug = format!("{user:?}");
748        let assert_common = |config: &WebSocketConfig| {
749            assert_eq!(config.headers, Vec::<(String, String)>::new());
750            assert_eq!(config.heartbeat_interval_secs, Some(10));
751            assert_eq!(config.connect_timeout_ms, Some(15_000));
752            assert_eq!(config.reconnect_delay_initial_ms, Some(250));
753            assert_eq!(config.reconnect_delay_max_ms, Some(5_000));
754            assert_eq!(config.reconnect_backoff_factor, Some(2.0));
755            assert_eq!(config.reconnect_jitter_ms, Some(200));
756            assert_eq!(config.reconnect_max_attempts, None);
757            // No data-silence timer: `PONG` arrives as a text frame and would
758            // refresh it, so liveness rests on the heartbeat timeout instead.
759            assert_eq!(config.idle_timeout_ms, None);
760            assert_eq!(config.backend, TransportBackend::Tungstenite);
761        };
762
763        assert_eq!(market.proxy_url.as_ref().unwrap().expose(), MARKET_PROXY);
764        assert_eq!(user.proxy_url.as_ref().unwrap().expose(), USER_PROXY);
765        assert_eq!(market_config.url, "ws://market.example/ws");
766        assert_eq!(user_config.url, "ws://user.example/ws");
767        assert_eq!(market_config.proxy_url.as_deref(), Some(MARKET_PROXY));
768        assert_eq!(user_config.proxy_url.as_deref(), Some(USER_PROXY));
769        assert_eq!(market_config.heartbeat_payload, None);
770        assert_eq!(user_config.heartbeat_payload.as_deref(), Some("PING"));
771        assert_common(&market_config);
772        assert_common(&user_config);
773        assert!(!market_debug.contains("market-proxy-secret"));
774        assert!(!user_debug.contains("user-proxy-secret"));
775    }
776}
nautilus_polymarket/websocket/client.rs

nautilus_polymarket/websocket/
client.rs
