Skip to main content

yapi_ai/auth/
xai.rs

1//! xAI (Grok/X subscription): an RFC 8628 device authorization against
2//! auth.x.ai. Port of `oauth/xai.ts`.
3
4use serde_json::{Map, Value};
5use tokio_util::sync::CancellationToken;
6use yapi_types::auth::OAuthCredential;
7
8use super::device::{DEVICE_CODE_GRANT, Poll, poll_device_code};
9use super::{
10    AuthError, AuthEvent, BoxFuture, Interaction, LoginOptions, OAuthProvider, now_ms, send,
11};
12
13const CLIENT_ID: &str = "b1a00492-073a-47ea-816f-4c329264a828";
14const SCOPE: &str = "openid profile email offline_access grok-cli:access api:access";
15/// Refresh this long before the reported expiry.
16const REFRESH_SKEW_MS: u64 = 5 * 60 * 1000;
17const DEFAULT_TOKEN_LIFETIME_SECONDS: f64 = 3600.0;
18
19/// The xAI sign-in and its endpoints.
20#[derive(Clone, Debug)]
21pub struct XaiOAuth {
22    /// Device authorization endpoint.
23    pub device_code_url: String,
24    /// Token endpoint.
25    pub token_url: String,
26}
27
28impl Default for XaiOAuth {
29    fn default() -> XaiOAuth {
30        XaiOAuth {
31            device_code_url: "https://auth.x.ai/oauth2/device/code".into(),
32            token_url: "https://auth.x.ai/oauth2/token".into(),
33        }
34    }
35}
36
37/// A response: success, status and JSON object body.
38struct Answer {
39    ok: bool,
40    status: u16,
41    body: Value,
42}
43
44/// POSTs `fields` as a form and returns the answer, with a non-object
45/// JSON body as `{}`.
46async fn request(
47    url: &str,
48    fields: &[(&str, &str)],
49    cancel: &CancellationToken,
50) -> Result<Answer, AuthError> {
51    let response = send(super::post_form(url, fields), cancel).await?;
52    let status = response.status().as_u16();
53    let bytes = response.bytes().await.unwrap_or_default();
54    let body = match serde_json::from_slice::<Value>(&bytes) {
55        Ok(object @ Value::Object(_)) => object,
56        Ok(_) => Value::Object(Map::new()),
57        Err(_) if cancel.is_cancelled() => return Err(AuthError::Cancelled),
58        Err(_) => {
59            return Err(AuthError::Failed(format!(
60                "xAI OAuth returned invalid JSON (HTTP {status})"
61            )));
62        }
63    };
64    Ok(Answer {
65        ok: (200..300).contains(&status),
66        status,
67        body,
68    })
69}
70
71fn failure(action: &str, answer: &Answer) -> String {
72    let detail: Vec<&str> = ["error", "error_description"]
73        .iter()
74        .filter_map(|key| answer.body[*key].as_str().filter(|value| !value.is_empty()))
75        .collect();
76    let detail = detail.join(": ");
77    let detail = if detail.is_empty() {
78        String::new()
79    } else {
80        format!(": {detail}")
81    };
82    format!("xAI OAuth {action} failed (HTTP {}){detail}", answer.status)
83}
84
85fn required<'a>(body: &'a Value, field: &str) -> Result<&'a str, AuthError> {
86    body[field]
87        .as_str()
88        .filter(|value| !value.is_empty())
89        .ok_or_else(|| AuthError::Failed(format!("Invalid xAI OAuth response field: {field}")))
90}
91
92fn positive(body: &Value, field: &str) -> Result<f64, AuthError> {
93    body[field]
94        .as_f64()
95        .filter(|value| value.is_finite() && *value > 0.0)
96        .ok_or_else(|| AuthError::Failed(format!("Invalid xAI OAuth response field: {field}")))
97}
98
99/// An https URL, the only kind opened in the browser.
100fn verification_uri(raw: &str) -> Result<String, AuthError> {
101    match url::Url::parse(raw) {
102        Ok(url) if url.scheme() == "https" => Ok(url.to_string()),
103        _ => Err(AuthError::failed(
104            "Untrusted verification URI in xAI OAuth response",
105        )),
106    }
107}
108
109fn credential(body: &Value, previous_refresh: Option<&str>) -> Result<OAuthCredential, AuthError> {
110    let access = required(body, "access_token")?.to_owned();
111    let refresh = match (&body["refresh_token"], previous_refresh) {
112        (Value::Null, Some(previous)) => previous.to_owned(),
113        _ => required(body, "refresh_token")?.to_owned(),
114    };
115    let expires_in = if body["expires_in"].is_null() {
116        DEFAULT_TOKEN_LIFETIME_SECONDS
117    } else {
118        positive(body, "expires_in")?
119    };
120    Ok(OAuthCredential {
121        access,
122        refresh,
123        expires: (now_ms() + (expires_in * 1000.0) as u64).saturating_sub(REFRESH_SKEW_MS) as f64,
124        extra: Map::new(),
125    })
126}
127
128impl XaiOAuth {
129    async fn login_xai(&self, interaction: &Interaction) -> Result<OAuthCredential, AuthError> {
130        let cancel = interaction.cancel();
131        let answer = request(
132            &self.device_code_url,
133            &[
134                ("client_id", CLIENT_ID),
135                ("scope", SCOPE),
136                ("referrer", "pi"),
137            ],
138            cancel,
139        )
140        .await?;
141        if !answer.ok {
142            return Err(AuthError::Failed(failure("device authorization", &answer)));
143        }
144        let body = &answer.body;
145        let interval = body["interval"]
146            .as_f64()
147            .filter(|value| value.is_finite() && *value > 0.0);
148        let complete = match body["verification_uri_complete"].as_str() {
149            Some(raw) if !raw.is_empty() => Some(verification_uri(raw)?),
150            _ => None,
151        };
152        let device_code = required(body, "device_code")?.to_owned();
153        let user_code = required(body, "user_code")?.to_owned();
154        let uri = verification_uri(required(body, "verification_uri")?)?;
155        let expires_in = positive(body, "expires_in")?;
156        interaction.notify(AuthEvent::DeviceCode {
157            user_code,
158            verification_uri: complete.unwrap_or(uri),
159            interval_seconds: interval,
160            expires_in_seconds: Some(expires_in),
161        });
162        poll_device_code(interval, Some(expires_in), true, cancel, || {
163            let device_code = device_code.clone();
164            async move {
165                let answer = request(
166                    &self.token_url,
167                    &[
168                        ("grant_type", DEVICE_CODE_GRANT),
169                        ("client_id", CLIENT_ID),
170                        ("device_code", &device_code),
171                    ],
172                    cancel,
173                )
174                .await?;
175                if answer.ok {
176                    return Ok(match credential(&answer.body, None) {
177                        Ok(credential) => Poll::Complete(credential),
178                        Err(err) => Poll::Failed(err.to_string()),
179                    });
180                }
181                Ok(match answer.body["error"].as_str() {
182                    Some("authorization_pending") => Poll::Pending,
183                    Some("slow_down") => Poll::SlowDown(answer.body["interval"].as_f64()),
184                    Some("access_denied" | "authorization_denied") => {
185                        Poll::Failed("xAI device authorization was denied".into())
186                    }
187                    Some("expired_token") => Poll::Failed("xAI device code expired".into()),
188                    _ => Poll::Failed(failure("device token polling", &answer)),
189                })
190            }
191        })
192        .await
193    }
194
195    async fn refresh_xai(
196        &self,
197        refresh: &str,
198        cancel: &CancellationToken,
199    ) -> Result<OAuthCredential, AuthError> {
200        let answer = request(
201            &self.token_url,
202            &[
203                ("grant_type", "refresh_token"),
204                ("client_id", CLIENT_ID),
205                ("refresh_token", refresh),
206            ],
207            cancel,
208        )
209        .await?;
210        if !answer.ok {
211            return Err(AuthError::Failed(failure("token refresh", &answer)));
212        }
213        credential(&answer.body, Some(refresh))
214    }
215}
216
217impl OAuthProvider for XaiOAuth {
218    fn name(&self) -> &str {
219        "xAI (Grok/X subscription)"
220    }
221
222    fn login_label(&self) -> Option<&str> {
223        Some("Sign in with SuperGrok or X Premium")
224    }
225
226    fn login<'a>(
227        &'a self,
228        interaction: &'a Interaction,
229        _options: &'a LoginOptions,
230    ) -> BoxFuture<'a, Result<OAuthCredential, AuthError>> {
231        Box::pin(self.login_xai(interaction))
232    }
233
234    fn refresh<'a>(
235        &'a self,
236        credential: &'a OAuthCredential,
237        cancel: &'a CancellationToken,
238    ) -> BoxFuture<'a, Result<OAuthCredential, AuthError>> {
239        Box::pin(self.refresh_xai(&credential.refresh, cancel))
240    }
241}