yapi_ai/auth/
kimi.rs

1//! Kimi Code (subscription): an RFC 8628 device authorization against
2//! auth.kimi.com, whose access token authenticates as a bearer header. Port
3//! of `oauth/kimi-coding.ts`.
4
5use std::time::Duration;
6
7use serde_json::{Map, Value};
8use tokio_util::sync::CancellationToken;
9use yapi_types::auth::OAuthCredential;
10
11use super::device::{DEVICE_CODE_GRANT, Poll, poll_device_code, positive, trusted_url};
12use super::{
13    AuthError, AuthEvent, BoxFuture, Interaction, LoginOptions, OAuthAuth, OAuthProvider,
14    json_body, now_ms, post_form, send,
15};
16
17const CLIENT_ID: &str = "17e5f671-d194-4dfb-9706-5516cb48c098";
18const DEFAULT_OAUTH_HOST: &str = "https://auth.kimi.com";
19const DEVICE_CODE_TIMEOUT_SECONDS: f64 = 15.0 * 60.0;
20const DEFAULT_POLL_INTERVAL_SECONDS: f64 = 5.0;
21const REFRESH_MAX_RETRIES: u32 = 3;
22const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
23
24/// The Kimi Code sign-in.
25#[derive(Clone, Debug, Default)]
26pub struct KimiOAuth {
27    /// Overrides the OAuth host, as `KIMI_CODE_OAUTH_HOST` does.
28    pub oauth_host: Option<String>,
29}
30
31impl KimiOAuth {
32    /// `KIMI_CODE_OAUTH_HOST`, else `KIMI_OAUTH_HOST`, else auth.kimi.com.
33    fn host(&self) -> String {
34        self.oauth_host
35            .clone()
36            .or_else(|| crate::credentials::provider_env_value("KIMI_CODE_OAUTH_HOST", None))
37            .or_else(|| crate::credentials::provider_env_value("KIMI_OAUTH_HOST", None))
38            .unwrap_or_else(|| DEFAULT_OAUTH_HOST.to_owned())
39            .trim_end_matches('/')
40            .to_owned()
41    }
42}
43
44/// A credential from a token response, or pi's error for `operation`.
45fn token(json: &Value, operation: &str) -> Result<OAuthCredential, AuthError> {
46    let access = json["access_token"]
47        .as_str()
48        .filter(|value| !value.is_empty());
49    let refresh = json["refresh_token"]
50        .as_str()
51        .filter(|value| !value.is_empty());
52    let expires_in = positive(&json["expires_in"]);
53    match (access, refresh, expires_in) {
54        (Some(access), Some(refresh), Some(expires_in)) => Ok(OAuthCredential {
55            access: access.to_owned(),
56            refresh: refresh.to_owned(),
57            expires: (now_ms() + (expires_in * 1000.0) as u64) as f64,
58            extra: Map::new(),
59        }),
60        _ => Err(AuthError::Failed(format!(
61            "Kimi Code token {operation} response missing fields: {}",
62            yapi_types::json::stringify(json)
63        ))),
64    }
65}
66
67impl KimiOAuth {
68    async fn login_kimi(&self, interaction: &Interaction) -> Result<OAuthCredential, AuthError> {
69        let host = self.host();
70        let cancel = interaction.cancel();
71        let response = send(
72            post_form(
73                format!("{host}/api/oauth/device_authorization"),
74                &[("client_id", CLIENT_ID)],
75            )
76            .timeout(REQUEST_TIMEOUT),
77            cancel,
78        )
79        .await?;
80        let status = response.status();
81        if !status.is_success() {
82            let text = response.text().await.unwrap_or_default();
83            let detail = if text.is_empty() {
84                String::new()
85            } else {
86                format!(": {text}")
87            };
88            return Err(AuthError::Failed(format!(
89                "Kimi Code device authorization failed with status {}{detail}",
90                status.as_u16()
91            )));
92        }
93        let json = json_body(response).await;
94        let (Some(device_code), Some(user_code), Some(_), Some(complete)) = (
95            json["device_code"].as_str(),
96            json["user_code"].as_str(),
97            trusted_url(&json["verification_uri"]),
98            trusted_url(&json["verification_uri_complete"]),
99        ) else {
100            return Err(AuthError::Failed(format!(
101                "Invalid Kimi Code device authorization response: {}",
102                yapi_types::json::stringify(&json)
103            )));
104        };
105        let interval = positive(&json["interval"]).unwrap_or(DEFAULT_POLL_INTERVAL_SECONDS);
106        let expires_in = positive(&json["expires_in"]).unwrap_or(DEVICE_CODE_TIMEOUT_SECONDS);
107        // pi opens the complete URI exactly as the server wrote it.
108        let complete = json["verification_uri_complete"]
109            .as_str()
110            .map_or(complete, str::to_owned);
111        interaction.notify(AuthEvent::DeviceCode {
112            user_code: user_code.to_owned(),
113            verification_uri: complete,
114            interval_seconds: Some(interval),
115            expires_in_seconds: Some(expires_in),
116        });
117        let token_url = format!("{host}/api/oauth/token");
118        poll_device_code(Some(interval), Some(expires_in), true, cancel, || {
119            let request = post_form(
120                &token_url,
121                &[
122                    ("client_id", CLIENT_ID),
123                    ("device_code", device_code),
124                    ("grant_type", DEVICE_CODE_GRANT),
125                ],
126            )
127            .timeout(REQUEST_TIMEOUT);
128            async move {
129                let response = send(request, cancel).await?;
130                let status = response.status().as_u16();
131                if status >= 500 {
132                    let text = response.text().await.unwrap_or_default();
133                    let detail = if text.is_empty() { String::new() } else { format!(": {text}") };
134                    return Ok(Poll::Failed(format!(
135                        "Kimi Code device token request failed with status {status}{detail}"
136                    )));
137                }
138                let json = json_body(response).await;
139                if (200..300).contains(&status) && json["access_token"].is_string() {
140                    return Ok(match token(&json, "poll") {
141                        Ok(credential) => Poll::Complete(credential),
142                        Err(err) => Poll::Failed(err.to_string()),
143                    });
144                }
145                let description = json["error_description"]
146                    .as_str()
147                    .map(|text| format!(": {text}"))
148                    .unwrap_or_default();
149                Ok(match json["error"].as_str() {
150                    Some("authorization_pending") => Poll::Pending,
151                    Some("slow_down") => Poll::SlowDown(positive(&json["interval"])),
152                    Some("expired_token") => Poll::Failed(
153                        "Kimi Code device authorization expired. Please restart login.".into(),
154                    ),
155                    Some("access_denied") => Poll::Failed("Kimi Code login was denied.".into()),
156                    Some(error) => Poll::Failed(format!(
157                        "Kimi Code device token request failed (status {status}): {error}{description}"
158                    )),
159                    None => Poll::Failed(format!(
160                        "Kimi Code device token request failed (status {status})"
161                    )),
162                })
163            }
164        })
165        .await
166    }
167
168    async fn refresh_kimi(
169        &self,
170        refresh: &str,
171        cancel: &CancellationToken,
172    ) -> Result<OAuthCredential, AuthError> {
173        let url = format!("{}/api/oauth/token", self.host());
174        let mut last = None;
175        for attempt in 0..=REFRESH_MAX_RETRIES {
176            if attempt > 0 {
177                let delay = Duration::from_millis(1000 << (attempt - 1));
178                cancel
179                    .run_until_cancelled(tokio::time::sleep(delay))
180                    .await
181                    .ok_or(AuthError::Cancelled)?;
182            }
183            if cancel.is_cancelled() {
184                return Err(AuthError::failed("Kimi Code token refresh aborted"));
185            }
186            let request = post_form(
187                &url,
188                &[
189                    ("client_id", CLIENT_ID),
190                    ("grant_type", "refresh_token"),
191                    ("refresh_token", refresh),
192                ],
193            )
194            .timeout(REQUEST_TIMEOUT);
195            let response = match send(request, cancel).await {
196                Ok(response) => response,
197                Err(AuthError::Cancelled) => return Err(AuthError::Cancelled),
198                Err(err) => {
199                    last = Some(err);
200                    continue;
201                }
202            };
203            let status = response.status().as_u16();
204            let json = json_body(response).await;
205            if (200..300).contains(&status) {
206                return token(&json, "refresh");
207            }
208            if status == 401 || status == 403 || json["error"] == "invalid_grant" {
209                let description = json["error_description"]
210                    .as_str()
211                    .map(|text| format!(": {text}"))
212                    .unwrap_or_default();
213                return Err(AuthError::Failed(format!(
214                    "Kimi Code token refresh unauthorized (status {status}){description}"
215                )));
216            }
217            if (status == 429 || status >= 500) && attempt < REFRESH_MAX_RETRIES {
218                last = Some(AuthError::Failed(format!(
219                    "Kimi Code token refresh failed with status {status}"
220                )));
221                continue;
222            }
223            return Err(AuthError::Failed(format!(
224                "Kimi Code token refresh failed with status {status}: {}",
225                yapi_types::json::stringify(&json)
226            )));
227        }
228        Err(last.unwrap_or_else(|| AuthError::failed("Kimi Code token refresh failed")))
229    }
230}
231
232impl OAuthProvider for KimiOAuth {
233    fn name(&self) -> &str {
234        "Kimi Code (subscription)"
235    }
236
237    fn login_label(&self) -> Option<&str> {
238        Some("Sign in with Kimi Code")
239    }
240
241    fn login<'a>(
242        &'a self,
243        interaction: &'a Interaction,
244        _options: &'a LoginOptions,
245    ) -> BoxFuture<'a, Result<OAuthCredential, AuthError>> {
246        Box::pin(self.login_kimi(interaction))
247    }
248
249    fn refresh<'a>(
250        &'a self,
251        credential: &'a OAuthCredential,
252        cancel: &'a CancellationToken,
253    ) -> BoxFuture<'a, Result<OAuthCredential, AuthError>> {
254        Box::pin(self.refresh_kimi(&credential.refresh, cancel))
255    }
256
257    fn to_auth<'a>(
258        &'a self,
259        credential: &'a OAuthCredential,
260    ) -> BoxFuture<'a, Result<OAuthAuth, AuthError>> {
261        let mut auth = OAuthAuth::default();
262        auth.headers.insert(
263            "Authorization".into(),
264            Some(format!("Bearer {}", credential.access)),
265        );
266        Box::pin(async move { Ok(auth) })
267    }
268}
yapi_ai/auth/kimi.rs

yapi_ai/auth/
kimi.rs
