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_common::live::get_runtime;
24use nautilus_network::{
25    mode::ConnectionMode,
26    websocket::{
27        AuthTracker, SubscriptionState, TransportBackend, WebSocketClient, WebSocketConfig,
28        channel_message_handler,
29    },
30};
31
32use super::{
33    handler::{FeedHandler, HandlerCommand},
34    messages::PolymarketWsMessage,
35};
36use crate::common::{
37    credential::Credential,
38    urls::{clob_ws_market_url, clob_ws_user_url},
39};
40
41const POLYMARKET_HEARTBEAT_SECS: u64 = 30;
42
43/// Polymarket WebSocket channel: market data or authenticated user data.
44#[derive(Debug, Clone, Copy, PartialEq, Eq)]
45pub enum WsChannel {
46    Market,
47    User,
48}
49
50// Market channel streams continuously; user channel can legitimately be quiet
51// when no orders or fills exist, so give it a longer window before treating
52// silence as a zombie connection.
53fn idle_timeout_ms_for(channel: WsChannel) -> u64 {
54    match channel {
55        WsChannel::Market => 60_000,
56        WsChannel::User => 300_000,
57    }
58}
59
60/// Lightweight handle for subscribing/unsubscribing to market data.
61///
62/// `Clone` + `Send` safe for use in spawned async tasks.
63#[derive(Clone, Debug)]
64pub struct WsSubscriptionHandle {
65    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
66}
67
68impl WsSubscriptionHandle {
69    /// Sends a market subscribe command to the handler.
70    pub async fn subscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
71        self.cmd_tx
72            .read()
73            .await
74            .send(HandlerCommand::SubscribeMarket(asset_ids))
75            .map_err(|e| anyhow::anyhow!("Failed to send SubscribeMarket: {e}"))
76    }
77
78    /// Sends a market unsubscribe command to the handler.
79    pub async fn unsubscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
80        self.cmd_tx
81            .read()
82            .await
83            .send(HandlerCommand::UnsubscribeMarket(asset_ids))
84            .map_err(|e| anyhow::anyhow!("Failed to send UnsubscribeMarket: {e}"))
85    }
86
87    // Constructs a handle around a raw command sender. Test-only: lets unit
88    // tests observe the commands the handle emits without spinning up the real
89    // feed handler.
90    #[cfg(test)]
91    pub(crate) fn from_sender(sender: tokio::sync::mpsc::UnboundedSender<HandlerCommand>) -> Self {
92        Self {
93            cmd_tx: Arc::new(tokio::sync::RwLock::new(sender)),
94        }
95    }
96}
97
98/// Provides a WebSocket client for the Polymarket CLOB API.
99///
100/// A single instance targets one channel (market or user). Use
101/// [`PolymarketWebSocketClient::new_market`] for public market data and
102/// [`PolymarketWebSocketClient::new_user`] for authenticated order/trade streams.
103#[derive(Debug)]
104pub struct PolymarketWebSocketClient {
105    channel: WsChannel,
106    url: String,
107    connection_mode: Arc<AtomicU8>,
108    signal: Arc<AtomicBool>,
109    cmd_tx: Arc<tokio::sync::RwLock<tokio::sync::mpsc::UnboundedSender<HandlerCommand>>>,
110    out_rx: Option<tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>>,
111    credential: Option<Credential>,
112    subscriptions: SubscriptionState,
113    auth_tracker: AuthTracker,
114    // Survives disconnect() so that connect() can replay a prior subscribe_user() call.
115    // Arc<AtomicBool> allows mutation from &self in subscribe_user().
116    user_subscribed: Arc<AtomicBool>,
117    task_handle: Option<tokio::task::JoinHandle<()>>,
118    subscribe_new_markets: bool,
119    transport_backend: TransportBackend,
120}
121
122impl PolymarketWebSocketClient {
123    /// Creates a new market-channel client (unauthenticated).
124    ///
125    /// If `base_url` is `None`, the default production URL is used.
126    #[must_use]
127    pub fn new_market(
128        base_url: Option<String>,
129        subscribe_new_markets: bool,
130        transport_backend: TransportBackend,
131    ) -> Self {
132        let url = base_url.unwrap_or_else(|| clob_ws_market_url().to_string());
133        Self::new_inner(
134            WsChannel::Market,
135            url,
136            None,
137            subscribe_new_markets,
138            transport_backend,
139        )
140    }
141
142    /// Creates a new user-channel client (authenticated).
143    ///
144    /// If `base_url` is `None`, the default production URL is used.
145    #[must_use]
146    pub fn new_user(
147        base_url: Option<String>,
148        credential: Credential,
149        transport_backend: TransportBackend,
150    ) -> Self {
151        let url = base_url.unwrap_or_else(|| clob_ws_user_url().to_string());
152        Self::new_inner(
153            WsChannel::User,
154            url,
155            Some(credential),
156            false,
157            transport_backend,
158        )
159    }
160
161    fn new_inner(
162        channel: WsChannel,
163        url: String,
164        credential: Option<Credential>,
165        subscribe_new_markets: bool,
166        transport_backend: TransportBackend,
167    ) -> Self {
168        let (placeholder_tx, _) = tokio::sync::mpsc::unbounded_channel();
169        Self {
170            channel,
171            url,
172            connection_mode: Arc::new(AtomicU8::new(ConnectionMode::Closed.as_u8())),
173            signal: Arc::new(AtomicBool::new(false)),
174            cmd_tx: Arc::new(tokio::sync::RwLock::new(placeholder_tx)),
175            out_rx: None,
176            credential,
177            subscriptions: SubscriptionState::new(':'),
178            auth_tracker: AuthTracker::new(),
179            user_subscribed: Arc::new(AtomicBool::new(false)),
180            task_handle: None,
181            subscribe_new_markets,
182            transport_backend,
183        }
184    }
185
186    /// Establishes the WebSocket connection and spawns the message handler.
187    ///
188    /// # Errors
189    ///
190    pub async fn connect(&mut self) -> anyhow::Result<()> {
191        let mode = ConnectionMode::from_atomic(&self.connection_mode);
192        if mode.is_active() || mode.is_reconnect() {
193            log::warn!("Polymarket WebSocket already connected or reconnecting");
194            return Ok(());
195        }
196
197        let (message_handler, raw_rx) = channel_message_handler();
198        let cfg = WebSocketConfig {
199            url: self.url.clone(),
200            headers: vec![],
201            heartbeat: Some(POLYMARKET_HEARTBEAT_SECS),
202            heartbeat_msg: None,
203            reconnect_timeout_ms: Some(15_000),
204            reconnect_delay_initial_ms: Some(250),
205            reconnect_delay_max_ms: Some(5_000),
206            reconnect_backoff_factor: Some(2.0),
207            reconnect_jitter_ms: Some(200),
208            reconnect_max_attempts: None,
209            idle_timeout_ms: Some(idle_timeout_ms_for(self.channel)),
210            backend: self.transport_backend,
211            proxy_url: None,
212        };
213
214        let client =
215            WebSocketClient::connect(cfg, Some(message_handler), None, None, vec![], None).await?;
216
217        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel::<HandlerCommand>();
218        let (out_tx, out_rx) = tokio::sync::mpsc::unbounded_channel::<PolymarketWsMessage>();
219
220        *self.cmd_tx.write().await = cmd_tx.clone();
221        self.out_rx = Some(out_rx);
222
223        let client_mode = client.connection_mode_atomic();
224        self.connection_mode = client_mode;
225
226        log::debug!("Polymarket WebSocket connected: {}", self.url);
227
228        cmd_tx
229            .send(HandlerCommand::SetClient(client))
230            .map_err(|e| anyhow::anyhow!("Failed to send SetClient: {e}"))?;
231
232        // Replay retained state onto the new session. Unlike the RECONNECTED sentinel
233        // path, a fresh connect() never fires resubscribe_all() inside the handler, so
234        // we must queue the commands here before the handler task is even spawned.
235        match self.channel {
236            WsChannel::Market => {
237                let topics = self.subscriptions.all_topics();
238                if !topics.is_empty() {
239                    log::debug!(
240                        "Replaying {} market subscription(s) onto new session",
241                        topics.len()
242                    );
243                    cmd_tx
244                        .send(HandlerCommand::SubscribeMarket(topics))
245                        .map_err(|e| anyhow::anyhow!("Failed to replay SubscribeMarket: {e}"))?;
246                }
247            }
248            WsChannel::User => {
249                if self.user_subscribed.load(Ordering::Relaxed) {
250                    log::debug!("Replaying user subscribe onto new session");
251                    cmd_tx
252                        .send(HandlerCommand::SubscribeUser)
253                        .map_err(|e| anyhow::anyhow!("Failed to replay SubscribeUser: {e}"))?;
254                }
255            }
256        }
257
258        let signal = Arc::clone(&self.signal);
259        let channel = self.channel;
260        let credential = self.credential.clone();
261        let subscriptions = self.subscriptions.clone();
262        let auth_tracker = self.auth_tracker.clone();
263        let user_subscribed = self.user_subscribed.load(Ordering::Relaxed);
264        let subscribe_new_markets = self.subscribe_new_markets;
265
266        let stream_handle = get_runtime().spawn(async move {
267            let mut handler = FeedHandler::new(
268                signal,
269                channel,
270                cmd_rx,
271                raw_rx,
272                out_tx,
273                credential,
274                subscriptions,
275                auth_tracker,
276                user_subscribed,
277                subscribe_new_markets,
278            );
279
280            loop {
281                match handler.next().await {
282                    Some(PolymarketWsMessage::Reconnected) => {
283                        log::info!("Polymarket WebSocket reconnected");
284                    }
285                    Some(msg) => {
286                        if handler.send(msg).is_err() {
287                            if handler.is_stopped() {
288                                log::debug!("Output channel closed, stopping handler");
289                            } else {
290                                log::error!("Output channel closed, stopping handler");
291                            }
292                            break;
293                        }
294                    }
295                    None => {
296                        if handler.is_stopped() {
297                            log::debug!("Stop signal received, ending handler task");
298                        } else {
299                            log::warn!("Polymarket WebSocket stream ended unexpectedly");
300                        }
301                        break;
302                    }
303                }
304            }
305            log::debug!("Polymarket WebSocket handler task completed");
306        });
307        self.task_handle = Some(stream_handle);
308        Ok(())
309    }
310
311    /// Force-close fallback for the sync `stop()` path.
312    /// Prefer `disconnect()` for graceful shutdown.
313    pub(crate) fn abort(&mut self) {
314        self.signal.store(true, Ordering::Relaxed);
315        self.connection_mode
316            .store(ConnectionMode::Closed.as_u8(), Ordering::SeqCst);
317
318        if let Some(handle) = self.task_handle.take() {
319            handle.abort();
320        }
321        self.auth_tracker.invalidate();
322    }
323
324    /// Disconnects the WebSocket connection.
325    pub async fn disconnect(&mut self) -> anyhow::Result<()> {
326        log::debug!("Disconnecting Polymarket WebSocket");
327        self.signal.store(true, Ordering::Relaxed);
328
329        if let Err(e) = self.cmd_tx.read().await.send(HandlerCommand::Disconnect) {
330            log::debug!("Failed to send disconnect (handler may already be shut down): {e}");
331        }
332
333        if let Some(handle) = self.task_handle.take() {
334            let abort_handle = handle.abort_handle();
335            tokio::select! {
336                result = handle => {
337                    match result {
338                        Ok(()) => log::debug!("Handler task completed"),
339                        Err(e) if e.is_cancelled() => log::debug!("Handler task was cancelled"),
340                        Err(e) => log::error!("Handler task error: {e:?}"),
341                    }
342                }
343                () = tokio::time::sleep(tokio::time::Duration::from_secs(2)) => {
344                    log::warn!("Timeout waiting for handler task, aborting");
345                    abort_handle.abort();
346                }
347            }
348        }
349        // Invalidate after the task has stopped so any in-flight auth_tracker.succeed()
350        // calls from the handler cannot race with and survive the invalidation.
351        self.auth_tracker.invalidate();
352        log::debug!("Polymarket WebSocket disconnected");
353        Ok(())
354    }
355
356    /// Returns `true` if the WebSocket is actively connected.
357    #[must_use]
358    pub fn is_active(&self) -> bool {
359        ConnectionMode::from_atomic(&self.connection_mode).is_active()
360    }
361
362    /// Returns the URL this client connects to.
363    #[must_use]
364    pub fn url(&self) -> &str {
365        &self.url
366    }
367
368    /// Returns the number of active market asset subscriptions (pending + confirmed).
369    #[must_use]
370    pub fn subscription_count(&self) -> usize {
371        self.subscriptions.all_topics().len()
372    }
373
374    /// Clears retained subscription/auth replay state.
375    ///
376    /// Useful for hard resets where the caller wants reconnect to start from a
377    /// clean slate rather than replaying a previous generation's topics.
378    pub(crate) fn clear_reconnect_state(&self) {
379        self.subscriptions.clear();
380        self.user_subscribed.store(false, Ordering::Relaxed);
381        self.auth_tracker.invalidate();
382    }
383
384    /// Returns `true` if the user channel has been authenticated.
385    #[must_use]
386    pub fn is_authenticated(&self) -> bool {
387        self.auth_tracker.is_authenticated()
388    }
389
390    /// Subscribe to market data for the given asset IDs.
391    ///
392    /// Sends a subscribe message immediately if connected; the IDs are also
393    /// retained so they are re-sent automatically on reconnect.
394    ///
395    /// # Errors
396    ///
397    /// Returns an error if called on a user-channel client (incompatible channel).
398    pub async fn subscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
399        if self.channel != WsChannel::Market {
400            anyhow::bail!(
401                "subscribe_market() requires a market-channel client (created with new_market())"
402            );
403        }
404        self.cmd_tx
405            .read()
406            .await
407            .send(HandlerCommand::SubscribeMarket(asset_ids))
408            .map_err(|e| anyhow::anyhow!("Failed to send SubscribeMarket: {e}"))
409    }
410
411    /// Remove asset IDs from the active subscription set.
412    ///
413    /// The IDs are dropped from the reconnect set so they will not be
414    /// re-subscribed after a reconnect. No wire message is sent.
415    ///
416    /// # Errors
417    ///
418    /// Returns an error if called on a user-channel client (incompatible channel).
419    pub async fn unsubscribe_market(&self, asset_ids: Vec<String>) -> anyhow::Result<()> {
420        if self.channel != WsChannel::Market {
421            anyhow::bail!(
422                "unsubscribe_market() requires a market-channel client (created with new_market())"
423            );
424        }
425        self.cmd_tx
426            .read()
427            .await
428            .send(HandlerCommand::UnsubscribeMarket(asset_ids))
429            .map_err(|e| anyhow::anyhow!("Failed to send UnsubscribeMarket: {e}"))
430    }
431
432    /// Authenticate and subscribe to the user channel.
433    ///
434    /// # Errors
435    ///
436    /// Returns an error if called on a market-channel client (no credentials available).
437    pub async fn subscribe_user(&self) -> anyhow::Result<()> {
438        if self.channel != WsChannel::User {
439            anyhow::bail!(
440                "subscribe_user() requires a user-channel client (created with new_user())"
441            );
442        }
443        self.cmd_tx
444            .read()
445            .await
446            .send(HandlerCommand::SubscribeUser)
447            .map_err(|e| anyhow::anyhow!("Failed to send SubscribeUser: {e}"))?;
448        // Set only after the command is successfully enqueued so a failed send does not
449        // leave user_subscribed=true and cause an unintended replay on the next connect().
450        self.user_subscribed.store(true, Ordering::Relaxed);
451        Ok(())
452    }
453
454    /// Returns a cloneable subscription handle for use in spawned tasks.
455    #[must_use]
456    pub fn clone_subscription_handle(&self) -> WsSubscriptionHandle {
457        WsSubscriptionHandle {
458            cmd_tx: Arc::clone(&self.cmd_tx),
459        }
460    }
461
462    /// Takes the message receiver, leaving `None` in its place.
463    ///
464    /// This is useful when the data client needs to spawn its own handler
465    /// task that reads messages independently of the WS client.
466    /// Subscription methods (`subscribe_market`, etc.) remain usable on `&self`.
467    #[must_use]
468    pub fn take_message_receiver(
469        &mut self,
470    ) -> Option<tokio::sync::mpsc::UnboundedReceiver<PolymarketWsMessage>> {
471        self.out_rx.take()
472    }
473
474    /// Receives the next message from the WebSocket handler.
475    ///
476    /// Returns `None` when the handler has disconnected or the receiver
477    /// was not yet initialized (call `connect` first).
478    pub async fn next_message(&mut self) -> Option<PolymarketWsMessage> {
479        if let Some(ref mut rx) = self.out_rx {
480            rx.recv().await
481        } else {
482            None
483        }
484    }
485}
486
487#[cfg(test)]
488mod tests {
489    use rstest::rstest;
490
491    use super::{WsChannel, idle_timeout_ms_for};
492
493    #[rstest]
494    #[case::market(WsChannel::Market, 60_000)]
495    #[case::user(WsChannel::User, 300_000)]
496    fn test_idle_timeout_ms_for_channel(#[case] channel: WsChannel, #[case] expected: u64) {
497        assert_eq!(idle_timeout_ms_for(channel), expected);
498    }
499}
nautilus_polymarket/websocket/client.rs

nautilus_polymarket/websocket/
client.rs
