Skip to main content

tradingview/live/
models.rs

1use core::fmt;
2
3use futures_util::stream::SplitStream;
4use serde::{Deserialize, Serialize};
5use serde_json::Value;
6use tokio::{net::TcpStream, sync::MutexGuard};
7use tokio_tungstenite::{
8    MaybeTlsStream, WebSocketStream,
9    tungstenite::{
10        http::{HeaderMap, HeaderValue},
11        protocol::Message,
12    },
13};
14use ustr::Ustr;
15
16use crate::{
17    Result, UA,
18    error::{Error, TradingViewError},
19    utils::format_packet,
20};
21use std::sync::LazyLock;
22
23pub static WEBSOCKET_HEADERS: LazyLock<HeaderMap<HeaderValue>> = LazyLock::new(|| {
24    let mut headers = HeaderMap::new();
25    headers.insert("Origin", "https://www.tradingview.com/".parse().unwrap());
26    headers.insert("User-Agent", UA.parse().unwrap());
27    headers
28});
29
30/// WebSocket event types dispatched by TradingView's data server.
31///
32/// Maps TradingView's wire protocol event names (`"timescale_update"`,
33/// `"du"`, `"qsd"`, etc.) to Rust enum variants.
34#[derive(Debug, Clone, PartialEq, Eq, Hash, Copy)]
35pub enum TradingViewDataEvent {
36    OnChartData,
37    OnChartDataUpdate,
38    OnQuoteData,
39    OnQuoteCompleted,
40    OnSeriesLoading,
41    OnSeriesCompleted,
42    OnSymbolResolved,
43    OnReplayOk,
44    OnReplayPoint,
45    OnReplayInstanceId,
46    OnReplayResolutions,
47    OnReplayDataEnd,
48    OnStudyLoading,
49    OnStudyCompleted,
50    OnError(TradingViewError),
51    UnknownEvent(Ustr),
52}
53
54impl From<String> for TradingViewDataEvent {
55    fn from(s: String) -> Self {
56        match s.as_str() {
57            "timescale_update" => TradingViewDataEvent::OnChartData,
58            "du" => TradingViewDataEvent::OnChartDataUpdate,
59
60            "qsd" => TradingViewDataEvent::OnQuoteData,
61            "quote_completed" => TradingViewDataEvent::OnQuoteCompleted,
62
63            "series_loading" => TradingViewDataEvent::OnSeriesLoading,
64            "series_completed" => TradingViewDataEvent::OnSeriesCompleted,
65
66            "symbol_resolved" => TradingViewDataEvent::OnSymbolResolved,
67
68            "replay_ok" => TradingViewDataEvent::OnReplayOk,
69            "replay_point" => TradingViewDataEvent::OnReplayPoint,
70            "replay_instance_id" => TradingViewDataEvent::OnReplayInstanceId,
71            "replay_resolutions" => TradingViewDataEvent::OnReplayResolutions,
72            "replay_data_end" => TradingViewDataEvent::OnReplayDataEnd,
73
74            "study_loading" => TradingViewDataEvent::OnStudyLoading,
75            "study_completed" => TradingViewDataEvent::OnStudyCompleted,
76
77            "symbol_error" => TradingViewDataEvent::OnError(TradingViewError::SymbolError),
78            "series_error" => TradingViewDataEvent::OnError(TradingViewError::SeriesError),
79            "critical_error" => TradingViewDataEvent::OnError(TradingViewError::CriticalError),
80            "study_error" => TradingViewDataEvent::OnError(TradingViewError::StudyError),
81            "protocol_error" => TradingViewDataEvent::OnError(TradingViewError::ProtocolError),
82            "replay_error" => TradingViewDataEvent::OnError(TradingViewError::ReplayError),
83
84            s => TradingViewDataEvent::UnknownEvent(s.into()),
85        }
86    }
87}
88
89impl From<Ustr> for TradingViewDataEvent {
90    fn from(s: Ustr) -> Self {
91        TradingViewDataEvent::from(s.to_string())
92    }
93}
94
95/// A serialized WebSocket message ready for transmission.
96#[derive(Debug, Clone, PartialEq, Serialize)]
97pub struct SocketMessageSer {
98    pub m: Value,
99    pub p: Value,
100}
101
102/// A deserialized WebSocket message from TradingView.
103#[derive(Debug, Clone, PartialEq, Deserialize)]
104pub struct SocketMessageDe {
105    pub m: Ustr,
106    pub p: Vec<Value>,
107    #[serde(default)]
108    pub t: u64, // Timestamp in seconds (0 when absent, e.g. error messages)
109    #[serde(default)]
110    pub t_ms: u64, // Timestamp in milliseconds (0 when absent)
111}
112
113impl SocketMessageSer {
114    pub fn new<M, P>(m: M, p: P) -> Self
115    where
116        M: Serialize,
117        P: Serialize,
118    {
119        let m = serde_json::to_value(m).expect("Failed to serialize Socket Message");
120        let p = serde_json::to_value(p).expect("Failed to serialize Socket Message");
121        SocketMessageSer { m, p }
122    }
123
124    pub fn to_message(&self) -> Result<Message> {
125        let msg = format_packet(self)?;
126        Ok(msg)
127    }
128}
129
130/// Server metadata sent by TradingView after a successful WebSocket connection.
131///
132/// Contains the session ID, server timestamp, and base URL for chart data.
133#[derive(Default, Debug, Clone, PartialEq, Serialize, Deserialize)]
134#[serde(rename_all = "camelCase")]
135pub struct SocketServerInfo {
136    #[serde(rename = "session_id")]
137    pub session_id: Ustr,
138    pub timestamp: i64,
139    pub timestamp_ms: i64,
140    pub release: Ustr,
141    #[serde(rename = "studies_metadata_hash")]
142    pub studies_metadata_hash: Ustr,
143    #[serde(rename = "auth_scheme_vsn")]
144    pub auth_scheme_vsn: i64,
145    pub protocol: Ustr,
146    pub via: Ustr,
147    #[serde(rename = "javastudies")]
148    pub sjavastudies: Vec<Ustr>,
149}
150
151impl fmt::Display for SocketServerInfo {
152    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
153        write!(
154            f,
155            "SocketServerInfo {{ session_id: {}, timestamp: {}, timestamp_ms: {}, release: {}, studies_metadata_hash: {}, auth_scheme_vsn: {}, protocol: {}, via: {}, sjavastudies: {:?} }}",
156            self.session_id,
157            self.timestamp,
158            self.timestamp_ms,
159            self.release,
160            self.studies_metadata_hash,
161            self.auth_scheme_vsn,
162            self.protocol,
163            self.via,
164            self.sjavastudies
165        )
166    }
167}
168
169#[derive(Debug, Clone, PartialEq, Deserialize, Serialize)]
170#[serde(untagged)]
171pub enum SocketMessage<T> {
172    SocketServerInfo(SocketServerInfo),
173    SocketMessage(T),
174    Heartbeat(u64),
175    Other(Value),
176    Unknown(String),
177}
178
179impl<T> SocketMessage<T> {
180    pub fn heartbeat_echo(&self) -> Option<String> {
181        match self {
182            SocketMessage::Heartbeat(counter) => {
183                let payload = format!("~h~{counter}");
184                Some(format!("~m~{}~m~{payload}", payload.len()))
185            }
186            _ => None,
187        }
188    }
189}
190
191/// Which TradingView data server tier to connect to.
192///
193/// `ProData` is the default and recommended server. `Data` and
194/// `DataExtended` are alternatives with different capabilities.
195#[derive(Default, Clone, Debug, PartialEq, Serialize, Deserialize, Copy, Eq)]
196pub enum DataServer {
197    #[default]
198    Data,
199    ProData,
200    WidgetData,
201    MobileData,
202}
203
204impl std::fmt::Display for DataServer {
205    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
206        match *self {
207            DataServer::Data => write!(f, "data"),
208            DataServer::ProData => write!(f, "prodata"),
209            DataServer::WidgetData => write!(f, "widgetdata"),
210            DataServer::MobileData => write!(f, "mobile-data"),
211        }
212    }
213}
214
215/// Trait for WebSocket message handling — serialize and deserialize from
216/// TradingView's wire format.
217///
218/// Implemented by [`SocketMessage`] for the two protocol variants.
219///
220/// [`SocketMessage`]: crate::live::models::SocketMessage
221pub trait Socket {
222    fn event_loop(
223        &self,
224        read: MutexGuard<SplitStream<WebSocketStream<MaybeTlsStream<TcpStream>>>>,
225    ) -> impl Future<Output = Result<()>> + Send;
226
227    fn handle_raw_messages(&self, raw: Message) -> impl Future<Output = Result<()>> + Send;
228
229    fn handle_parsed_messages(
230        &self,
231        messages: Vec<SocketMessage<SocketMessageDe>>,
232        raw: &Message,
233    ) -> impl Future<Output = Result<()>> + Send;
234
235    fn handle_message_data(
236        &self,
237        message: SocketMessageDe,
238    ) -> impl Future<Output = Result<()>> + Send;
239
240    fn handle_error(&self, error: Error, context: Ustr) -> impl Future<Output = Result<()>> + Send;
241}
242
243// ---------------------------------------------------------------------------
244// Tests
245// ---------------------------------------------------------------------------
246
247#[cfg(test)]
248mod tests {
249    use super::*;
250
251    #[test]
252    fn test_study_loading_maps_to_on_study_loading() {
253        let event = TradingViewDataEvent::from("study_loading".to_string());
254        assert_eq!(event, TradingViewDataEvent::OnStudyLoading);
255    }
256
257    #[test]
258    fn test_series_loading_is_not_study_loading() {
259        let event = TradingViewDataEvent::from("series_loading".to_string());
260        assert_eq!(event, TradingViewDataEvent::OnSeriesLoading);
261    }
262
263    #[test]
264    fn test_study_completed_maps_correctly() {
265        let event = TradingViewDataEvent::from("study_completed".to_string());
266        assert_eq!(event, TradingViewDataEvent::OnStudyCompleted);
267    }
268
269    #[test]
270    fn test_series_completed_maps_correctly() {
271        let event = TradingViewDataEvent::from("series_completed".to_string());
272        assert_eq!(event, TradingViewDataEvent::OnSeriesCompleted);
273    }
274
275    #[test]
276    fn test_all_study_and_series_events_are_distinct() {
277        let loading = TradingViewDataEvent::from("study_loading".to_string());
278        let completed = TradingViewDataEvent::from("study_completed".to_string());
279        let series_loading = TradingViewDataEvent::from("series_loading".to_string());
280        let series_completed = TradingViewDataEvent::from("series_completed".to_string());
281
282        // All four should be distinct
283        assert_ne!(loading, series_loading);
284        assert_ne!(completed, series_completed);
285        assert_ne!(loading, completed);
286        assert_ne!(series_loading, series_completed);
287
288        // Verify mapping expectations
289        assert_eq!(loading, TradingViewDataEvent::OnStudyLoading);
290        assert_eq!(completed, TradingViewDataEvent::OnStudyCompleted);
291        assert_eq!(series_loading, TradingViewDataEvent::OnSeriesLoading);
292        assert_eq!(series_completed, TradingViewDataEvent::OnSeriesCompleted);
293    }
294
295    #[test]
296    fn test_ustr_from_maps_correctly() {
297        let event: TradingViewDataEvent = ustr::ustr("study_loading").into();
298        assert_eq!(event, TradingViewDataEvent::OnStudyLoading);
299    }
300}