nautilus_polymarket/websocket/
handler.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 message handler for the Polymarket CLOB API.
17
18use std::{
19    sync::{
20        Arc,
21        atomic::{AtomicBool, Ordering},
22    },
23    time::Duration,
24};
25
26use ahash::AHashMap;
27use nautilus_network::{
28    RECONNECTED,
29    websocket::{AuthTracker, SubscriptionState, WebSocketClient},
30};
31use serde_json::value::RawValue;
32use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender}; // tokio-import-ok
33use tokio_tungstenite::tungstenite::Message;
34use ustr::Ustr;
35
36use super::{
37    client::{POLYMARKET_HEARTBEAT_PAYLOAD, POLYMARKET_HEARTBEAT_SECS, WsChannel},
38    messages::{
39        MarketInitialSubscribeRequest, MarketSubscribeRequest, MarketUnsubscribeRequest,
40        MarketWsMessage, PolymarketWsAuth, PolymarketWsMessage, UserSubscribeRequest,
41        UserWsMessage,
42    },
43};
44use crate::{common::credential::Credential, http::error::sanitize_error_text};
45
46const INITIAL_DUMP: bool = true;
47
48/// Commands sent from the outer client to the inner message handler.
49#[derive(Debug)]
50pub enum HandlerCommand {
51    /// Set the WebSocketClient for the handler to use.
52    SetClient(WebSocketClient),
53    /// Disconnect the WebSocket connection.
54    Disconnect,
55    /// Add asset IDs to the market-channel subscription set and send a subscribe message.
56    SubscribeMarket(Vec<String>),
57    /// Remove asset IDs from the subscription set (no wire message needed).
58    UnsubscribeMarket(Vec<String>),
59    /// Send the authenticated subscribe message on the user channel.
60    SubscribeUser,
61}
62
63pub(super) struct FeedHandler {
64    signal: Arc<AtomicBool>,
65    channel: WsChannel,
66    client: Option<WebSocketClient>,
67    cmd_rx: UnboundedReceiver<HandlerCommand>,
68    raw_rx: UnboundedReceiver<(u64, Message)>,
69    out_tx: UnboundedSender<PolymarketWsMessage>,
70    credential: Option<Credential>,
71    subscriptions: SubscriptionState,
72    discovery_subscribed: Arc<AtomicBool>,
73    initial_market_replay: Option<(Vec<String>, u64)>,
74    auth_tracker: AuthTracker,
75    // True once SubscribeUser has been explicitly requested by the caller
76    user_subscribed: bool,
77    // True once the current market-channel session has sent its initial subscribe payload.
78    market_subscription_initialized: bool,
79    market_heartbeat_next: Option<(tokio::time::Instant, u64)>,
80    // Assets awaiting first authoritative data, keyed by the connection that wrote the subscribe.
81    market_subscription_epochs: AHashMap<String, u64>,
82    // Overflow buffer for batched frames, drained before reading the next raw message
83    message_buffer: Vec<PolymarketWsMessage>,
84    // Whether to include `custom_feature_enabled: true` in the initial subscribe
85    subscribe_new_markets: bool,
86}
87
88impl FeedHandler {
89    #[expect(clippy::too_many_arguments)]
90    pub(super) fn new(
91        signal: Arc<AtomicBool>,
92        channel: WsChannel,
93        client: Option<WebSocketClient>,
94        cmd_rx: UnboundedReceiver<HandlerCommand>,
95        raw_rx: UnboundedReceiver<(u64, Message)>,
96        out_tx: UnboundedSender<PolymarketWsMessage>,
97        credential: Option<Credential>,
98        subscriptions: SubscriptionState,
99        discovery_subscribed: Arc<AtomicBool>,
100        initial_market_replay: Option<(Vec<String>, u64)>,
101        auth_tracker: AuthTracker,
102        user_subscribed: bool,
103        subscribe_new_markets: bool,
104    ) -> Self {
105        Self {
106            signal,
107            channel,
108            client,
109            cmd_rx,
110            raw_rx,
111            out_tx,
112            credential,
113            subscriptions,
114            discovery_subscribed,
115            initial_market_replay,
116            auth_tracker,
117            user_subscribed,
118            market_subscription_initialized: false,
119            market_heartbeat_next: None,
120            market_subscription_epochs: AHashMap::new(),
121            message_buffer: Vec::new(),
122            subscribe_new_markets,
123        }
124    }
125
126    pub(super) fn send(&self, msg: PolymarketWsMessage) -> Result<(), String> {
127        self.out_tx
128            .send(msg)
129            .map_err(|e| format!("Failed to send message: {e}"))
130    }
131
132    pub(super) fn is_stopped(&self) -> bool {
133        self.signal.load(Ordering::Relaxed)
134    }
135
136    async fn send_subscribe_market(&mut self, asset_ids: &[String], connection_epoch: Option<u64>) {
137        let Some(ref client) = self.client else {
138            log::warn!("No client available for market subscribe");
139            return;
140        };
141
142        let connection_epoch = connection_epoch.unwrap_or_else(|| client.connection_epoch());
143
144        for id in asset_ids {
145            self.market_subscription_epochs.remove(id);
146            self.subscriptions.mark_subscribe(id);
147        }
148
149        let payload = if self.market_subscription_initialized {
150            serde_json::to_string(&MarketSubscribeRequest {
151                assets_ids: asset_ids.to_vec(),
152                operation: "subscribe",
153                initial_dump: INITIAL_DUMP,
154                custom_feature_enabled: self.subscribe_new_markets,
155            })
156        } else {
157            serde_json::to_string(&MarketInitialSubscribeRequest {
158                assets_ids: asset_ids.to_vec(),
159                msg_type: "market",
160                initial_dump: INITIAL_DUMP,
161                custom_feature_enabled: self.subscribe_new_markets,
162            })
163        };
164
165        match payload {
166            Ok(payload) => {
167                let result = client
168                    .send_text_on_connection(payload, None, connection_epoch)
169                    .await;
170
171                if let Err(e) = result {
172                    for id in asset_ids {
173                        self.market_subscription_epochs.remove(id);
174                        self.subscriptions.mark_failure(id);
175                    }
176                    log::error!("Failed to send market subscribe: {e}");
177                } else {
178                    for id in asset_ids {
179                        if self.market_subscription_pending(id) {
180                            self.market_subscription_epochs
181                                .insert(id.clone(), connection_epoch);
182                        }
183                    }
184
185                    if !self.market_subscription_initialized {
186                        self.market_subscription_initialized = true;
187                        self.schedule_market_heartbeat(connection_epoch);
188                    }
189                }
190            }
191            Err(e) => {
192                for id in asset_ids {
193                    self.market_subscription_epochs.remove(id);
194                    self.subscriptions.mark_failure(id);
195                }
196                log::error!("Failed to serialize market subscribe request: {e}");
197            }
198        }
199    }
200
201    async fn send_unsubscribe_market(&self, asset_ids: &[String]) {
202        let Some(ref client) = self.client else {
203            log::warn!("No client available for market unsubscribe");
204            return;
205        };
206
207        let req = MarketUnsubscribeRequest {
208            assets_ids: asset_ids.to_vec(),
209            operation: "unsubscribe",
210        };
211
212        match serde_json::to_string(&req) {
213            Ok(payload) => {
214                if let Err(e) = client.send_text(payload, None).await {
215                    log::error!("Failed to send market unsubscribe: {e}");
216                }
217            }
218            Err(e) => log::error!("Failed to serialize market unsubscribe request: {e}"),
219        }
220    }
221
222    async fn send_subscribe_user(&self) {
223        let Some(ref client) = self.client else {
224            log::warn!("No client available for user subscribe");
225            return;
226        };
227        let Some(cred) = &self.credential else {
228            log::error!("User channel subscribe requires credential");
229            return;
230        };
231
232        let req = UserSubscribeRequest {
233            auth: PolymarketWsAuth {
234                api_key: cred.api_key().to_string(),
235                secret: cred.api_secret(),
236                passphrase: cred.passphrase().to_string(),
237            },
238            msg_type: "user",
239        };
240
241        // Begin auth tracking; discard receiver, state is queried via is_authenticated()
242        drop(self.auth_tracker.begin());
243
244        match serde_json::to_string(&req) {
245            Ok(payload) => {
246                // auth_tracker.succeed() is NOT called here; sending the request only
247                // confirms delivery to the server, not that the credentials were accepted.
248                // succeed() is called in next() when the server actually sends user-channel
249                // data, which is the real confirmation that authentication worked.
250                if let Err(e) = client.send_text(payload, None).await {
251                    self.auth_tracker.fail(e.to_string());
252                    log::error!("Failed to send user subscribe: {e}");
253                }
254            }
255            Err(e) => {
256                self.auth_tracker.fail(format!("Serialize error: {e}"));
257                log::error!("Failed to serialize user subscribe request: {e}");
258            }
259        }
260    }
261
262    async fn resubscribe_all(&mut self, connection_epoch: u64) {
263        match self.channel {
264            WsChannel::Market => {
265                let ids = self.subscriptions.reset_after_reconnect();
266                if ids.is_empty() && !self.discovery_subscribed.load(Ordering::Relaxed) {
267                    return;
268                }
269                log::info!(
270                    "Restoring market subscription state after reconnect: assets={}, discovery={}",
271                    ids.len(),
272                    self.discovery_subscribed.load(Ordering::Relaxed),
273                );
274                self.send_subscribe_market(&ids, Some(connection_epoch))
275                    .await;
276            }
277            WsChannel::User => {
278                if self.user_subscribed {
279                    log::info!("Re-authenticating user channel after reconnect");
280                    self.send_subscribe_user().await;
281                }
282            }
283        }
284    }
285
286    fn parse_messages(&self, text: &str) -> Vec<PolymarketWsMessage> {
287        // When `subscribe_new_markets` is enabled, Polymarket's WSS periodically
288        // sends the plain-text string "NO NEW ASSETS" as a heartbeat/ack.
289        if text == "NO NEW ASSETS" {
290            return vec![];
291        }
292
293        // Reply to the application-level `PING` heartbeat, which is not JSON
294        if text == "PONG" {
295            return vec![];
296        }
297
298        match self.channel {
299            WsChannel::Market => {
300                if let Ok(msgs) = serde_json::from_str::<Vec<&RawValue>>(text) {
301                    msgs.into_iter()
302                        .filter_map(|raw| match MarketWsMessage::parse(raw.get()) {
303                            Ok(msg) => Some(PolymarketWsMessage::Market(msg)),
304                            Err(e) => {
305                                log::warn!("Failed to parse market WS batch element: {e}");
306                                None
307                            }
308                        })
309                        .collect()
310                } else {
311                    match MarketWsMessage::parse(text) {
312                        Ok(msg) => vec![PolymarketWsMessage::Market(msg)],
313                        Err(e) => {
314                            log::warn!(
315                                "Failed to parse market WS message: {e}; payload={}",
316                                sanitize_error_text(text)
317                            );
318                            vec![]
319                        }
320                    }
321                }
322            }
323            WsChannel::User => {
324                if let Ok(msgs) = UserWsMessage::parse_batch(text) {
325                    msgs.into_iter().map(PolymarketWsMessage::User).collect()
326                } else {
327                    match UserWsMessage::parse(text) {
328                        Ok(msg) => vec![PolymarketWsMessage::User(msg)],
329                        Err(e) => {
330                            log::warn!(
331                                "Failed to parse user WS message: {e}; payload={}",
332                                sanitize_error_text(text)
333                            );
334                            vec![]
335                        }
336                    }
337                }
338            }
339        }
340    }
341
342    pub(super) async fn next(&mut self) -> Option<PolymarketWsMessage> {
343        if !self.message_buffer.is_empty() {
344            return Some(self.message_buffer.remove(0));
345        }
346
347        if let Some((asset_ids, connection_epoch)) = self.initial_market_replay.take() {
348            self.send_subscribe_market(&asset_ids, Some(connection_epoch))
349                .await;
350        }
351
352        loop {
353            let market_heartbeat_next = self.market_heartbeat_next;
354
355            tokio::select! {
356                connection_epoch = async move {
357                    if let Some((deadline, connection_epoch)) = market_heartbeat_next {
358                        tokio::time::sleep_until(deadline).await;
359                        connection_epoch
360                    } else {
361                        std::future::pending::<u64>().await
362                    }
363                } => {
364                    self.send_market_heartbeat(connection_epoch).await;
365                    self.schedule_market_heartbeat(connection_epoch);
366                }
367                Some(cmd) = self.cmd_rx.recv() => {
368                    match cmd {
369                        HandlerCommand::SetClient(client) => {
370                            log::debug!("Setting WebSocket client in handler");
371                            self.client = Some(client);
372                        }
373                        HandlerCommand::Disconnect => {
374                            log::debug!("Handler received disconnect command");
375
376                            if let Some(ref client) = self.client {
377                                client.disconnect().await;
378                            }
379                            self.signal.store(true, Ordering::SeqCst);
380                            return None;
381                        }
382                        HandlerCommand::SubscribeMarket(ids) => {
383                            if self.subscribe_new_markets && ids.is_empty() {
384                                self.discovery_subscribed.store(true, Ordering::Relaxed);
385                            }
386                            self.send_subscribe_market(&ids, None).await;
387                        }
388                        HandlerCommand::UnsubscribeMarket(ids) => {
389                            for id in &ids {
390                                self.market_subscription_epochs.remove(id);
391                                self.subscriptions.mark_unsubscribe(id);
392                            }
393                            self.send_unsubscribe_market(&ids).await;
394                            for id in &ids {
395                                self.subscriptions.confirm_unsubscribe(id);
396                            }
397                        }
398                        HandlerCommand::SubscribeUser => {
399                            self.user_subscribed = true;
400                            self.send_subscribe_user().await;
401                        }
402                    }
403                }
404                Some((connection_epoch, raw)) = self.raw_rx.recv() => {
405                    match raw {
406                        Message::Text(text) => {
407                            if text == RECONNECTED {
408                                self.market_subscription_initialized = false;
409                                self.market_heartbeat_next = None;
410                                self.resubscribe_all(connection_epoch).await;
411                                return Some(PolymarketWsMessage::Reconnected);
412                            }
413                            let msgs = self.parse_messages(&text);
414                            if msgs.is_empty() {
415                                continue;
416                            }
417
418                            if self.channel == WsChannel::Market {
419                                self.confirm_market_subscriptions(connection_epoch, &msgs);
420                            } else {
421                                // Receiving any user-channel data confirms the server accepted the
422                                // credentials; mark auth as successful on the first delivery.
423                                self.auth_tracker.succeed();
424                            }
425                            // Buffer msgs[1..] so they are returned in order on subsequent
426                            // next() calls; returning first directly preserves 0,1,2,...,n order
427                            let mut iter = msgs.into_iter();
428                            let first = iter.next().unwrap();
429                            self.message_buffer.extend(iter);
430                            return Some(first);
431                        }
432                        Message::Ping(data) => {
433                            if let Some(ref client) = self.client
434                                && let Err(e) = client.send_pong(data.to_vec()).await
435                            {
436                                log::warn!("Failed to send pong: {e}");
437                            }
438                        }
439                        Message::Close(_) => {
440                            log::debug!("WebSocket close frame received");
441                            return None;
442                        }
443                        _ => {}
444                    }
445                }
446                else => return None,
447            }
448        }
449    }
450
451    fn schedule_market_heartbeat(&mut self, connection_epoch: u64) {
452        self.market_heartbeat_next = Some((
453            tokio::time::Instant::now() + Duration::from_secs(POLYMARKET_HEARTBEAT_SECS),
454            connection_epoch,
455        ));
456    }
457
458    async fn send_market_heartbeat(&self, connection_epoch: u64) {
459        let Some(ref client) = self.client else {
460            return;
461        };
462
463        if let Err(e) = client
464            .send_text_on_connection(
465                POLYMARKET_HEARTBEAT_PAYLOAD.to_string(),
466                None,
467                connection_epoch,
468            )
469            .await
470        {
471            log::debug!("Failed to send market heartbeat: {e}");
472        }
473    }
474
475    fn confirm_market_subscriptions(
476        &mut self,
477        connection_epoch: u64,
478        messages: &[PolymarketWsMessage],
479    ) {
480        for message in messages {
481            let asset_id = match message {
482                PolymarketWsMessage::Market(MarketWsMessage::Book(book)) => &book.asset_id,
483                PolymarketWsMessage::Market(MarketWsMessage::LastTradePrice(trade)) => {
484                    &trade.asset_id
485                }
486                _ => continue,
487            };
488
489            let was_sent_on_connection = self
490                .market_subscription_epochs
491                .get(asset_id.as_str())
492                .is_some_and(|epoch| *epoch == connection_epoch);
493
494            if was_sent_on_connection && self.market_subscription_pending(asset_id.as_str()) {
495                self.subscriptions.confirm_subscribe(asset_id.as_str());
496                self.market_subscription_epochs.remove(asset_id.as_str());
497            }
498        }
499    }
500
501    fn market_subscription_pending(&self, asset_id: &str) -> bool {
502        let channel_level = Ustr::from("");
503        let asset_id = Ustr::from(asset_id);
504        self.subscriptions
505            .pending_subscribe()
506            .get(&asset_id)
507            .is_some_and(|symbols| symbols.contains(&channel_level))
508    }
509}
510
511#[cfg(test)]
512mod tests {
513    use std::time::Duration;
514
515    use futures_util::StreamExt;
516    use nautilus_common::testing::wait_until_async;
517    use nautilus_network::websocket::{TransportBackend, WebSocketConfig, channel_message_handler};
518    use parking_lot::Mutex;
519    use rstest::{fixture, rstest};
520    use serde_json::{Value, json};
521
522    use super::*;
523    use crate::common::enums::PolymarketOrderSide;
524
525    const MARKET_ASSET_ID: &str =
526        "71321045679252212594626385532706912750332728571942532289631379312455583992563";
527
528    #[fixture]
529    fn market_handler() -> FeedHandler {
530        feed_handler(WsChannel::Market)
531    }
532
533    #[fixture]
534    fn user_handler() -> FeedHandler {
535        feed_handler(WsChannel::User)
536    }
537
538    fn feed_handler(channel: WsChannel) -> FeedHandler {
539        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
540        let (_raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
541        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
542
543        FeedHandler::new(
544            Arc::new(AtomicBool::new(false)),
545            channel,
546            None,
547            cmd_rx,
548            raw_rx,
549            out_tx,
550            None,
551            SubscriptionState::new(':'),
552            Arc::new(AtomicBool::new(false)),
553            None,
554            AuthTracker::new(),
555            false,
556            false,
557        )
558    }
559
560    async fn recording_server() -> (String, Arc<Mutex<Vec<String>>>) {
561        let listener = tokio::net::TcpListener::bind("127.0.0.1:0")
562            .await
563            .expect("bind recording server");
564        let addr = listener.local_addr().expect("recording server address");
565        let messages = Arc::new(Mutex::new(Vec::new()));
566        let received = Arc::clone(&messages);
567
568        tokio::spawn(async move {
569            let (stream, _) = listener.accept().await.expect("accept websocket client");
570            let mut socket = tokio_tungstenite::accept_async(stream)
571                .await
572                .expect("accept websocket handshake");
573
574            while let Some(message) = socket.next().await {
575                match message.expect("read websocket message") {
576                    Message::Text(text) => received.lock().push(text.to_string()),
577                    Message::Close(_) => break,
578                    _ => {}
579                }
580            }
581        });
582
583        (format!("ws://{addr}"), messages)
584    }
585
586    async fn recording_client(url: String) -> WebSocketClient {
587        let config = WebSocketConfig::builder()
588            .url(url)
589            .backend(TransportBackend::Tungstenite)
590            .build()
591            .expect("valid websocket config");
592        let (message_handler, _message_rx) = channel_message_handler();
593        WebSocketClient::builder()
594            .config(config)
595            .message_handler(message_handler)
596            .connect()
597            .await
598            .expect("connect websocket client")
599    }
600
601    fn market_handler_with(
602        client: WebSocketClient,
603    ) -> (
604        FeedHandler,
605        UnboundedSender<(u64, Message)>,
606        SubscriptionState,
607    ) {
608        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
609        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
610        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
611        let subscriptions = SubscriptionState::new(':');
612        let handler = FeedHandler::new(
613            Arc::new(AtomicBool::new(false)),
614            WsChannel::Market,
615            Some(client),
616            cmd_rx,
617            raw_rx,
618            out_tx,
619            None,
620            subscriptions.clone(),
621            Arc::new(AtomicBool::new(false)),
622            None,
623            AuthTracker::new(),
624            false,
625            false,
626        );
627
628        (handler, raw_tx, subscriptions)
629    }
630
631    #[rstest]
632    #[tokio::test]
633    async fn initial_market_replay_recovers_on_current_connection_epoch() {
634        let (url, messages) = recording_server().await;
635        let client = recording_client(url).await;
636
637        let (_cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
638        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
639        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
640        let mut handler = FeedHandler::new(
641            Arc::new(AtomicBool::new(false)),
642            WsChannel::Market,
643            Some(client),
644            cmd_rx,
645            raw_rx,
646            out_tx,
647            None,
648            SubscriptionState::new(':'),
649            Arc::new(AtomicBool::new(true)),
650            Some((vec![], 1)),
651            AuthTracker::new(),
652            false,
653            true,
654        );
655        raw_tx
656            .send((0, Message::Text(RECONNECTED.into())))
657            .expect("queue reconnect notification");
658
659        assert!(matches!(
660            handler.next().await,
661            Some(PolymarketWsMessage::Reconnected),
662        ));
663        handler
664            .client
665            .as_ref()
666            .expect("websocket client")
667            .send_text_on_connection("barrier".to_string(), None, 0)
668            .await
669            .expect("send barrier on current connection");
670
671        wait_until_async(
672            || {
673                let messages = Arc::clone(&messages);
674                async move { messages.lock().len() >= 2 }
675            },
676            Duration::from_secs(1),
677        )
678        .await;
679
680        {
681            let messages = messages.lock();
682            assert_eq!(messages.len(), 2);
683            assert_eq!(
684                serde_json::from_str::<Value>(&messages[0]).expect("valid subscribe payload"),
685                json!({
686                    "assets_ids": [],
687                    "type": "market",
688                    "initial_dump": true,
689                    "custom_feature_enabled": true,
690                }),
691            );
692            assert_eq!(messages[1], "barrier");
693        }
694
695        handler
696            .client
697            .as_ref()
698            .expect("websocket client")
699            .disconnect()
700            .await;
701    }
702
703    #[rstest]
704    #[tokio::test(start_paused = true)]
705    async fn market_text_heartbeat_follows_initial_subscription() {
706        let (url, messages) = recording_server().await;
707        let client = recording_client(url).await;
708        let connection_epoch = client.connection_epoch();
709        let (cmd_tx, cmd_rx) = tokio::sync::mpsc::unbounded_channel();
710        let (raw_tx, raw_rx) = tokio::sync::mpsc::unbounded_channel();
711        let (out_tx, _out_rx) = tokio::sync::mpsc::unbounded_channel();
712        let mut handler = FeedHandler::new(
713            Arc::new(AtomicBool::new(false)),
714            WsChannel::Market,
715            Some(client),
716            cmd_rx,
717            raw_rx,
718            out_tx,
719            None,
720            SubscriptionState::new(':'),
721            Arc::new(AtomicBool::new(false)),
722            None,
723            AuthTracker::new(),
724            false,
725            false,
726        );
727
728        handler
729            .send_subscribe_market(&[MARKET_ASSET_ID.to_string()], None)
730            .await;
731        wait_for_recorded_messages(&messages, 1).await;
732
733        let task = tokio::spawn(async move {
734            let message = handler.next().await;
735            (handler, message)
736        });
737        tokio::task::yield_now().await;
738        tokio::time::advance(Duration::from_secs(POLYMARKET_HEARTBEAT_SECS)).await;
739        wait_for_recorded_messages(&messages, 2).await;
740
741        raw_tx
742            .send((connection_epoch, Message::Text(RECONNECTED.into())))
743            .expect("queue reconnect notification");
744        let (mut handler, message) = task.await.expect("join handler task");
745        assert!(matches!(message, Some(PolymarketWsMessage::Reconnected)));
746        wait_for_recorded_messages(&messages, 3).await;
747
748        let task = tokio::spawn(async move {
749            let message = handler.next().await;
750            (handler, message)
751        });
752        tokio::task::yield_now().await;
753        tokio::time::advance(Duration::from_secs(POLYMARKET_HEARTBEAT_SECS)).await;
754        wait_for_recorded_messages(&messages, 4).await;
755
756        cmd_tx
757            .send(HandlerCommand::Disconnect)
758            .expect("queue disconnect");
759        let (_, message) = task.await.expect("join handler task");
760        assert!(message.is_none());
761
762        let messages = messages.lock().clone();
763        let expected_subscription = json!({
764            "assets_ids": [MARKET_ASSET_ID],
765            "type": "market",
766            "initial_dump": true,
767        });
768        assert_eq!(messages.len(), 4);
769        assert_eq!(
770            serde_json::from_str::<Value>(&messages[0]).expect("valid subscribe payload"),
771            expected_subscription,
772        );
773        assert_eq!(messages[1], POLYMARKET_HEARTBEAT_PAYLOAD);
774        assert_eq!(
775            serde_json::from_str::<Value>(&messages[2]).expect("valid replay payload"),
776            expected_subscription,
777        );
778        assert_eq!(messages[3], POLYMARKET_HEARTBEAT_PAYLOAD);
779    }
780
781    #[rstest]
782    #[tokio::test]
783    async fn market_heartbeat_stays_bound_to_subscribed_connection() {
784        let (url, messages) = recording_server().await;
785        let client = recording_client(url).await;
786        let connection_epoch = client.connection_epoch();
787        let connection_epoch_atomic = client.connection_epoch_atomic();
788        let (mut handler, raw_tx, _) = market_handler_with(client);
789
790        handler
791            .send_subscribe_market(&[MARKET_ASSET_ID.to_string()], None)
792            .await;
793        wait_for_recorded_messages(&messages, 1).await;
794
795        let replacement_epoch = connection_epoch + 1;
796        connection_epoch_atomic.store(replacement_epoch, Ordering::Release);
797        handler.send_market_heartbeat(connection_epoch).await;
798        raw_tx
799            .send((replacement_epoch, Message::Text(RECONNECTED.into())))
800            .expect("queue reconnect notification");
801
802        assert!(matches!(
803            handler.next().await,
804            Some(PolymarketWsMessage::Reconnected),
805        ));
806        handler
807            .client
808            .as_ref()
809            .expect("websocket client")
810            .send_text_on_connection("barrier".to_string(), None, replacement_epoch)
811            .await
812            .expect("send barrier on replacement connection");
813        wait_until_async(
814            || {
815                let messages = Arc::clone(&messages);
816                async move {
817                    messages
818                        .lock()
819                        .last()
820                        .is_some_and(|message| message == "barrier")
821                }
822            },
823            Duration::from_secs(1),
824        )
825        .await;
826
827        let messages = messages.lock().clone();
828        let expected_subscription = json!({
829            "assets_ids": [MARKET_ASSET_ID],
830            "type": "market",
831            "initial_dump": true,
832        });
833        assert_eq!(messages.len(), 3);
834        assert_eq!(
835            serde_json::from_str::<Value>(&messages[0]).expect("valid subscribe payload"),
836            expected_subscription,
837        );
838        assert_eq!(
839            serde_json::from_str::<Value>(&messages[1]).expect("valid replay payload"),
840            expected_subscription,
841        );
842        assert_eq!(messages[2], "barrier");
843
844        handler
845            .client
846            .as_ref()
847            .expect("websocket client")
848            .disconnect()
849            .await;
850    }
851
852    async fn wait_for_recorded_messages(messages: &Arc<Mutex<Vec<String>>>, expected: usize) {
853        wait_until_async(
854            || {
855                let messages = Arc::clone(messages);
856                async move { messages.lock().len() == expected }
857            },
858            Duration::from_secs(1),
859        )
860        .await;
861    }
862
863    #[rstest]
864    #[case(include_str!("../../test_data/ws_market_book_msg.json"))]
865    #[case(include_str!("../../test_data/ws_market_last_trade_msg.json"))]
866    #[tokio::test]
867    async fn market_subscription_confirms_from_first_book_or_trade(#[case] payload: &str) {
868        let (url, messages) = recording_server().await;
869        let client = recording_client(url).await;
870        let (mut handler, raw_tx, subscriptions) = market_handler_with(client);
871
872        handler
873            .send_subscribe_market(&[MARKET_ASSET_ID.to_string()], None)
874            .await;
875        wait_until_async(
876            || {
877                let messages = Arc::clone(&messages);
878                async move { !messages.lock().is_empty() }
879            },
880            Duration::from_secs(1),
881        )
882        .await;
883
884        assert_eq!(
885            subscriptions.pending_subscribe_topics(),
886            vec![MARKET_ASSET_ID]
887        );
888        assert_eq!(subscriptions.len(), 0);
889
890        raw_tx
891            .send((0, Message::Text(payload.into())))
892            .expect("queue market data");
893        assert!(matches!(
894            handler.next().await,
895            Some(PolymarketWsMessage::Market(_)),
896        ));
897
898        assert!(subscriptions.pending_subscribe_topics().is_empty());
899        assert_eq!(subscriptions.len(), 1);
900
901        handler
902            .client
903            .as_ref()
904            .expect("websocket client")
905            .disconnect()
906            .await;
907    }
908
909    #[rstest]
910    fn unsolicited_market_data_does_not_create_subscription(mut market_handler: FeedHandler) {
911        let messages =
912            market_handler.parse_messages(include_str!("../../test_data/ws_market_book_msg.json"));
913
914        market_handler.confirm_market_subscriptions(0, &messages);
915
916        assert!(market_handler.subscriptions.is_empty());
917    }
918
919    #[rstest]
920    fn market_batch_confirms_trade_but_not_price_change(mut market_handler: FeedHandler) {
921        let price_change_asset_id = "101";
922        let trade_asset_id = "202";
923        market_handler
924            .subscriptions
925            .mark_subscribe(price_change_asset_id);
926        market_handler.subscriptions.mark_subscribe(trade_asset_id);
927        market_handler
928            .market_subscription_epochs
929            .insert(price_change_asset_id.to_string(), 0);
930        market_handler
931            .market_subscription_epochs
932            .insert(trade_asset_id.to_string(), 0);
933        let messages = market_handler.parse_messages(include_str!(
934            "../../test_data/ws_market_mixed_known_unknown.json"
935        ));
936
937        market_handler.confirm_market_subscriptions(0, &messages);
938
939        assert_eq!(
940            market_handler.subscriptions.pending_subscribe_topics(),
941            vec![price_change_asset_id]
942        );
943        assert_eq!(market_handler.subscriptions.len(), 1);
944    }
945
946    #[rstest]
947    fn market_subscription_confirmation_requires_sent_current_epoch(
948        mut market_handler: FeedHandler,
949    ) {
950        market_handler.subscriptions.mark_subscribe(MARKET_ASSET_ID);
951        let messages =
952            market_handler.parse_messages(include_str!("../../test_data/ws_market_book_msg.json"));
953
954        market_handler.confirm_market_subscriptions(0, &messages);
955        assert_eq!(
956            market_handler.subscriptions.pending_subscribe_topics(),
957            vec![MARKET_ASSET_ID]
958        );
959        assert_eq!(market_handler.subscriptions.len(), 0);
960
961        market_handler
962            .market_subscription_epochs
963            .insert(MARKET_ASSET_ID.to_string(), 1);
964        market_handler.confirm_market_subscriptions(0, &messages);
965        assert_eq!(
966            market_handler.subscriptions.pending_subscribe_topics(),
967            vec![MARKET_ASSET_ID]
968        );
969        assert_eq!(market_handler.subscriptions.len(), 0);
970
971        market_handler.confirm_market_subscriptions(1, &messages);
972        assert!(
973            market_handler
974                .subscriptions
975                .pending_subscribe_topics()
976                .is_empty()
977        );
978        assert_eq!(market_handler.subscriptions.len(), 1);
979    }
980
981    #[rstest]
982    #[tokio::test]
983    async fn reconnect_replay_requires_current_connection_data() {
984        let cancelled_asset_id = "cancelled-asset";
985        let (url, _) = recording_server().await;
986        let client = recording_client(url).await;
987        let connection_epoch = client.connection_epoch();
988        let (mut handler, raw_tx, subscriptions) = market_handler_with(client);
989        subscriptions.mark_subscribe(MARKET_ASSET_ID);
990        subscriptions.confirm_subscribe(MARKET_ASSET_ID);
991        subscriptions.mark_subscribe(cancelled_asset_id);
992        subscriptions.confirm_subscribe(cancelled_asset_id);
993        subscriptions.mark_unsubscribe(cancelled_asset_id);
994
995        handler.resubscribe_all(connection_epoch).await;
996
997        assert_eq!(
998            subscriptions.pending_subscribe_topics(),
999            vec![MARKET_ASSET_ID]
1000        );
1001        assert!(subscriptions.pending_unsubscribe_topics().is_empty());
1002        assert_eq!(subscriptions.len(), 0);
1003
1004        raw_tx
1005            .send((
1006                connection_epoch,
1007                Message::Text(include_str!("../../test_data/ws_market_book_msg.json").into()),
1008            ))
1009            .expect("queue market data");
1010        assert!(matches!(
1011            handler.next().await,
1012            Some(PolymarketWsMessage::Market(_)),
1013        ));
1014
1015        assert!(subscriptions.pending_subscribe_topics().is_empty());
1016        assert_eq!(subscriptions.len(), 1);
1017
1018        handler
1019            .client
1020            .as_ref()
1021            .expect("websocket client")
1022            .disconnect()
1023            .await;
1024    }
1025
1026    #[rstest]
1027    #[tokio::test]
1028    async fn failed_market_subscribe_stays_pending_for_reconnect_replay() {
1029        let (url, _) = recording_server().await;
1030        let client = recording_client(url).await;
1031        client.disconnect().await;
1032        let (mut handler, raw_tx, subscriptions) = market_handler_with(client);
1033        subscriptions.mark_subscribe(MARKET_ASSET_ID);
1034        subscriptions.confirm_subscribe(MARKET_ASSET_ID);
1035
1036        assert!(subscriptions.pending_subscribe_topics().is_empty());
1037        assert_eq!(subscriptions.len(), 1);
1038
1039        handler
1040            .send_subscribe_market(&[MARKET_ASSET_ID.to_string()], None)
1041            .await;
1042
1043        assert_eq!(
1044            subscriptions.pending_subscribe_topics(),
1045            vec![MARKET_ASSET_ID]
1046        );
1047        assert_eq!(subscriptions.len(), 0);
1048
1049        raw_tx
1050            .send((
1051                0,
1052                Message::Text(include_str!("../../test_data/ws_market_book_msg.json").into()),
1053            ))
1054            .expect("queue stale market data");
1055        assert!(matches!(
1056            handler.next().await,
1057            Some(PolymarketWsMessage::Market(_)),
1058        ));
1059        assert_eq!(
1060            subscriptions.pending_subscribe_topics(),
1061            vec![MARKET_ASSET_ID]
1062        );
1063        assert_eq!(subscriptions.len(), 0);
1064
1065        let (replay_url, messages) = recording_server().await;
1066        let replay_client = recording_client(replay_url).await;
1067        let connection_epoch = replay_client.connection_epoch();
1068        handler.client = Some(replay_client);
1069        handler.resubscribe_all(connection_epoch).await;
1070        wait_until_async(
1071            || {
1072                let messages = Arc::clone(&messages);
1073                async move { !messages.lock().is_empty() }
1074            },
1075            Duration::from_secs(1),
1076        )
1077        .await;
1078
1079        {
1080            let messages = messages.lock();
1081            assert_eq!(messages.len(), 1);
1082            assert_eq!(
1083                serde_json::from_str::<Value>(&messages[0]).expect("valid subscribe payload"),
1084                json!({
1085                    "assets_ids": [MARKET_ASSET_ID],
1086                    "type": "market",
1087                    "initial_dump": true,
1088                }),
1089            );
1090        }
1091        assert_eq!(
1092            subscriptions.pending_subscribe_topics(),
1093            vec![MARKET_ASSET_ID]
1094        );
1095        assert_eq!(subscriptions.len(), 0);
1096
1097        handler
1098            .client
1099            .as_ref()
1100            .expect("websocket client")
1101            .disconnect()
1102            .await;
1103    }
1104
1105    #[rstest]
1106    fn test_parse_market_batch_skips_unknown_event(market_handler: FeedHandler) {
1107        let messages = market_handler.parse_messages(include_str!(
1108            "../../test_data/ws_market_mixed_known_unknown.json"
1109        ));
1110
1111        assert_eq!(messages.len(), 2);
1112
1113        let PolymarketWsMessage::Market(MarketWsMessage::PriceChange(quotes)) = &messages[0] else {
1114            panic!("Expected first message to be a price change");
1115        };
1116        assert_eq!(
1117            quotes.market.as_str(),
1118            "0x1111111111111111111111111111111111111111111111111111111111111111"
1119        );
1120        assert_eq!(quotes.timestamp, "1700000000001");
1121        assert_eq!(quotes.price_changes.len(), 1);
1122
1123        let quote = &quotes.price_changes[0];
1124        assert_eq!(quote.asset_id.as_str(), "101");
1125        assert_eq!(quote.price, "0.37");
1126        assert_eq!(quote.side, PolymarketOrderSide::Buy);
1127        assert_eq!(quote.size, "12.5");
1128        assert_eq!(quote.hash, "price-change-hash");
1129        assert_eq!(quote.best_bid.as_deref(), Some("0.36"));
1130        assert_eq!(quote.best_ask.as_deref(), Some("0.38"));
1131
1132        let PolymarketWsMessage::Market(MarketWsMessage::LastTradePrice(trade)) = &messages[1]
1133        else {
1134            panic!("Expected second message to be a last trade price");
1135        };
1136        assert_eq!(
1137            trade.market.as_str(),
1138            "0x2222222222222222222222222222222222222222222222222222222222222222"
1139        );
1140        assert_eq!(trade.asset_id.as_str(), "202");
1141        assert_eq!(trade.fee_rate_bps, "17");
1142        assert_eq!(trade.price, "0.63");
1143        assert_eq!(trade.side, PolymarketOrderSide::Sell);
1144        assert_eq!(trade.size, "4.25");
1145        assert_eq!(trade.timestamp, "1700000000003");
1146        assert_eq!(trade.transaction_hash.as_deref(), Some("0xtrade-hash"));
1147    }
1148
1149    #[rstest]
1150    fn test_parse_market_single_message(market_handler: FeedHandler) {
1151        let messages = market_handler.parse_messages(include_str!(
1152            "../../test_data/ws_market_last_trade_msg.json"
1153        ));
1154
1155        assert_eq!(messages.len(), 1);
1156
1157        let PolymarketWsMessage::Market(MarketWsMessage::LastTradePrice(trade)) = &messages[0]
1158        else {
1159            panic!("Expected a last trade price");
1160        };
1161        assert_eq!(
1162            trade.market.as_str(),
1163            "0xdd22472e552920b8438158ea7238bfadfa4f736aa4cee91a6b86c39ead110917"
1164        );
1165        assert_eq!(
1166            trade.asset_id.as_str(),
1167            "71321045679252212594626385532706912750332728571942532289631379312455583992563"
1168        );
1169        assert_eq!(trade.fee_rate_bps, "0");
1170        assert_eq!(trade.price, "0.51");
1171        assert_eq!(trade.side, PolymarketOrderSide::Buy);
1172        assert_eq!(trade.size, "25.0");
1173        assert_eq!(trade.timestamp, "1703875202000");
1174        assert!(trade.transaction_hash.is_none());
1175    }
1176
1177    #[rstest]
1178    fn test_parse_user_batch(user_handler: FeedHandler) {
1179        let messages =
1180            user_handler.parse_messages(include_str!("../../test_data/ws_user_batch_msg.json"));
1181        let actual: Vec<UserWsMessage> = messages
1182            .into_iter()
1183            .map(|message| match message {
1184                PolymarketWsMessage::User(message) => message,
1185                other => panic!("Expected user message, received {other:?}"),
1186            })
1187            .collect();
1188        let expected: Vec<UserWsMessage> =
1189            serde_json::from_str(include_str!("../../test_data/ws_user_batch_msg.json"))
1190                .expect("user batch fixture should deserialize");
1191
1192        assert_eq!(actual, expected);
1193    }
1194}
nautilus_polymarket/websocket/handler.rs

nautilus_polymarket/websocket/
handler.rs
