Skip to main content

zai_rs/realtime/
client.rs

1//! [`RealtimeClient`] — entry point for the realtime API.
2
3use std::sync::Arc;
4
5use super::session::SessionBuilder;
6use crate::{
7    ZaiResult,
8    client::{ApiFamily, endpoint::EndpointConfig, secret::ApiSecret},
9};
10
11/// Authentication mode for the realtime WebSocket handshake.
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub enum AuthMode {
14    /// Server-side Bearer auth: `Authorization: Bearer {API_KEY}` (default).
15    Bearer,
16    /// JWT auth using a short-lived token derived locally from the API key.
17    /// Only the derived token is sent in the WebSocket handshake. Keeping the
18    /// original key away from an untrusted browser or device remains the
19    /// application's responsibility.
20    Jwt {
21        /// Token validity in seconds.
22        ttl_seconds: i64,
23    },
24}
25
26/// Entry point for the realtime API.
27///
28/// Construct with [`RealtimeClient::new`], optionally switch to JWT auth via
29/// [`RealtimeClient::with_jwt`], then start a session with
30/// [`RealtimeClient::session`].
31///
32/// ```rust,no_run
33/// use zai_rs::{model::GLM_realtime_flash, realtime::RealtimeClient};
34///
35/// # async fn go(key: String) -> zai_rs::ZaiResult<()> {
36/// let session = RealtimeClient::new(key)
37///     .session(GLM_realtime_flash {})
38///     .build()
39///     .await?;
40/// # Ok(())
41/// # }
42/// ```
43#[derive(Clone)]
44pub struct RealtimeClient {
45    api_key: Arc<ApiSecret>,
46    auth: AuthMode,
47    endpoint_config: EndpointConfig,
48}
49
50impl RealtimeClient {
51    /// Create a realtime client using Bearer auth and the default endpoints.
52    pub fn new(api_key: impl Into<String>) -> Self {
53        Self {
54            api_key: Arc::new(ApiSecret::new(api_key)),
55            auth: AuthMode::Bearer,
56            endpoint_config: EndpointConfig::defaults()
57                .expect("built-in realtime endpoint must be valid"),
58        }
59    }
60
61    /// Switch to JWT auth, signing tokens valid for `ttl_seconds`.
62    ///
63    /// The value is validated when the session is built and must be between one
64    /// second and seven days, inclusive.
65    pub fn with_jwt(mut self, ttl_seconds: i64) -> Self {
66        self.auth = AuthMode::Jwt { ttl_seconds };
67        self
68    }
69
70    /// Switch (back) to server-side Bearer auth.
71    pub fn with_bearer(mut self) -> Self {
72        self.auth = AuthMode::Bearer;
73        self
74    }
75
76    /// Override the full endpoint config (base URLs).
77    pub fn with_endpoint_config(mut self, config: EndpointConfig) -> Self {
78        self.endpoint_config = config;
79        self
80    }
81
82    /// Override only the realtime base URL (`wss://...`).
83    ///
84    /// Only the realtime family is replaced; all other endpoint overrides are
85    /// preserved. Invalid or insecure public URLs return an error.
86    pub fn try_with_realtime_base(mut self, base: impl AsRef<str>) -> ZaiResult<Self> {
87        self.endpoint_config =
88            self.endpoint_config
89                .with_base(ApiFamily::Realtime, base.as_ref(), false)?;
90        Ok(self)
91    }
92
93    /// Begin building a realtime session for `model`.
94    ///
95    /// The model bound is checked at compile time via [`super::RealtimeModel`].
96    pub fn session<M: super::RealtimeModel>(&self, model: M) -> SessionBuilder {
97        let realtime_url = self.realtime_url();
98        let model_name: String = model.into();
99        SessionBuilder::new(
100            Arc::clone(&self.api_key),
101            self.auth,
102            realtime_url,
103            model_name,
104        )
105    }
106
107    /// The resolved realtime WebSocket URL.
108    pub fn realtime_url(&self) -> String {
109        self.endpoint_config.base(ApiFamily::Realtime).to_string()
110    }
111
112    /// Current auth mode.
113    pub fn auth(&self) -> &AuthMode {
114        &self.auth
115    }
116}
117
118impl std::fmt::Debug for RealtimeClient {
119    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
120        formatter
121            .debug_struct("RealtimeClient")
122            .field("api_key", &self.api_key)
123            .field("auth", &self.auth)
124            .field("endpoint_config", &self.endpoint_config)
125            .finish()
126    }
127}
128
129#[cfg(test)]
130mod tests {
131    use super::*;
132
133    #[test]
134    fn debug_redacts_api_key() {
135        let client = RealtimeClient::new("abcdefghij.0123456789abcdef");
136        let debug = format!("{client:?}");
137        assert!(debug.contains("[REDACTED]"));
138        assert!(!debug.contains("abcdefghij"));
139    }
140
141    #[test]
142    fn realtime_override_accepts_an_owned_non_static_base() {
143        let base = format!("wss://{}/realtime", "example.com");
144        let client = RealtimeClient::new("abcdefghij.0123456789abcdef")
145            .try_with_realtime_base(&base)
146            .unwrap();
147        assert_eq!(client.realtime_url(), "wss://example.com/realtime");
148        assert!(
149            client
150                .try_with_realtime_base("ws://public.example/realtime")
151                .is_err()
152        );
153    }
154}