Skip to main content

mj_controller/
hel_quota.rs

1//! One-pane quota collection for Mjolnir harness profiles.
2
3use std::collections::{BTreeMap, BTreeSet, HashMap};
4use std::path::Path;
5use std::sync::{Arc, Mutex};
6use std::time::{Duration, SystemTime};
7
8use anyhow::{Context, Result, bail};
9use chrono::{DateTime, Datelike, Days, FixedOffset, Local, NaiveDate, NaiveTime, TimeZone};
10use serde::{Deserialize, Serialize};
11use serde_json::Value;
12
13use crate::claude_usage;
14use crate::codex_usage::{self, CodexUsageClient, CodexUsageStatus};
15use crate::grok_usage;
16use hel::hel_config::HarnessKind;
17use hel::hel_config::harness_authentication_marker;
18use hel::hel_credentials::{
19    MAX_CREDENTIAL_BYTES, credential_expiry, credential_fingerprint, credential_freshness,
20};
21
22pub const API_LABEL: &str = "API";
23
24#[derive(Debug, Clone)]
25pub struct QuotaRefreshRequest {
26    pub profile_id: String,
27    pub harness: HarnessKind,
28    pub source_home: std::path::PathBuf,
29    pub environment: BTreeMap<String, String>,
30    pub cwd: std::path::PathBuf,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
34pub struct QuotaWindow {
35    pub label: String,
36    pub remaining_percent: Option<u8>,
37    pub used: Option<i64>,
38    pub limit: Option<i64>,
39    pub resets: Option<String>,
40    #[serde(default)]
41    pub resets_at_epoch_seconds: Option<i64>,
42}
43
44#[derive(Debug, Clone, Serialize, Deserialize, PartialEq, Eq)]
45pub struct ProfileQuota {
46    pub profile_id: String,
47    pub harness: HarnessKind,
48    pub windows: Vec<QuotaWindow>,
49    pub extra: Option<String>,
50    pub error: Option<String>,
51    pub refreshed_at_epoch_seconds: u64,
52}
53
54#[derive(Debug, Clone, PartialEq, Eq)]
55pub struct QuotaRefreshOutcome {
56    pub report: ProfileQuota,
57    pub credentials_changed: bool,
58}
59
60impl ProfileQuota {
61    pub fn weekly_window(&self) -> Option<&QuotaWindow> {
62        self.windows
63            .iter()
64            .find(|window| is_weekly_quota_window(&window.label))
65    }
66
67    pub fn five_hour_window(&self) -> Option<&QuotaWindow> {
68        self.windows
69            .iter()
70            .find(|window| is_short_quota_window(&window.label))
71    }
72
73    /// Whether the report says the profile is usage-priced: an API-billed
74    /// harness has no subscription window to fill, so it reports the API label
75    /// in place of one rather than inventing a percentage.
76    pub fn is_usage_priced(&self) -> bool {
77        self.error.is_none() && self.windows.is_empty() && self.extra.as_deref() == Some(API_LABEL)
78    }
79
80    pub fn five_hour_projects_exhaustion(&self) -> bool {
81        self.five_hour_window().is_some_and(|window| {
82            projects_exhaustion_before_reset(window, self.refreshed_at_epoch_seconds)
83        })
84    }
85
86    pub fn compact(&self) -> String {
87        if let Some(error) = &self.error {
88            return quota_error_label(error);
89        }
90        let mut seen_resets = BTreeSet::new();
91        let mut parts = self
92            .windows
93            .iter()
94            .filter(|window| {
95                !is_short_quota_window(&window.label)
96                    || projects_exhaustion_before_reset(window, self.refreshed_at_epoch_seconds)
97            })
98            .map(|window| {
99                let usage = match (window.remaining_percent, window.used, window.limit) {
100                    (Some(remaining), _, _) => format!("{remaining}% left"),
101                    (_, Some(used), Some(limit)) => format!("{used}/{limit}"),
102                    _ => "available".to_string(),
103                };
104                match window
105                    .resets
106                    .as_ref()
107                    .filter(|reset| seen_resets.insert((*reset).clone()))
108                {
109                    Some(reset) => format!("{} {usage}, resets {reset}", window.label),
110                    None => format!("{} {usage}", window.label),
111                }
112            })
113            .collect::<Vec<_>>();
114        if let Some(extra) = &self.extra {
115            parts.push(extra.clone());
116        }
117        if parts.is_empty() {
118            "no quota windows reported".to_string()
119        } else {
120            parts.join(" ยท ")
121        }
122    }
123
124    pub fn error_label(&self) -> Option<String> {
125        self.error.as_deref().map(quota_error_label)
126    }
127}
128
129fn quota_error_label(error: &str) -> String {
130    if error == claude_usage::LOGIN_EXPIRED {
131        claude_usage::LOGIN_EXPIRED.to_string()
132    } else {
133        format!("unavailable: {error}")
134    }
135}
136
137/// The dashboard's long-window column. A harness billed monthly rather than
138/// weekly belongs in the same column; the label itself names the real period.
139fn is_weekly_quota_window(label: &str) -> bool {
140    matches!(
141        label.to_ascii_lowercase().as_str(),
142        "week" | "weekly" | "7d" | "month" | "monthly"
143    )
144}
145
146fn is_short_quota_window(label: &str) -> bool {
147    matches!(
148        label.to_ascii_lowercase().as_str(),
149        "5h" | "5-hour" | "5 hour"
150    )
151}
152
153/// Whether this window is on course to run out before it resets.
154///
155/// Published so the phone can show the warning without recomputing a rule that
156/// belongs here, beside the data it reads.
157pub fn projects_exhaustion(window: &QuotaWindow, now: u64) -> bool {
158    projects_exhaustion_before_reset(window, now)
159}
160
161fn projects_exhaustion_before_reset(window: &QuotaWindow, now: u64) -> bool {
162    const FIVE_HOURS_SECONDS: i64 = 5 * 60 * 60;
163    let Some(reset) = window.resets_at_epoch_seconds else {
164        return false;
165    };
166    let Ok(now) = i64::try_from(now) else {
167        return false;
168    };
169    let remaining_time = reset - now;
170    let elapsed = FIVE_HOURS_SECONDS - remaining_time;
171    if remaining_time <= 0 || elapsed <= 0 || elapsed >= FIVE_HOURS_SECONDS {
172        return false;
173    }
174    if let (Some(used), Some(limit)) = (window.used, window.limit)
175        && limit > 0
176    {
177        return i128::from(used.clamp(0, limit)) * i128::from(FIVE_HOURS_SECONDS)
178            > i128::from(limit) * i128::from(elapsed);
179    }
180    window
181        .remaining_percent
182        .is_some_and(|remaining| i64::from(100 - remaining) * FIVE_HOURS_SECONDS > 100 * elapsed)
183}
184
185#[derive(Default)]
186pub struct QuotaManager {
187    codex_clients: HashMap<String, CodexUsageClient>,
188    reports: BTreeMap<String, ProfileQuota>,
189}
190
191impl QuotaManager {
192    pub fn reports(&self) -> &BTreeMap<String, ProfileQuota> {
193        &self.reports
194    }
195
196    /// Refresh each profile independently so one slow harness cannot delay the
197    /// others. `on_report` runs per profile in completion order, so fast
198    /// harnesses report without waiting for the slowest one in the batch.
199    pub async fn refresh_profiles<F, Fut>(
200        &mut self,
201        requests: Vec<QuotaRefreshRequest>,
202        mut on_report: F,
203    ) where
204        F: FnMut(QuotaRefreshOutcome) -> Fut,
205        Fut: Future<Output = ()> + Send,
206    {
207        let batch = requests
208            .iter()
209            .map(|request| request.profile_id.clone())
210            .collect::<BTreeSet<_>>();
211        let mut tasks = tokio::task::JoinSet::new();
212        for request in requests {
213            let client = self.codex_clients.remove(&request.profile_id);
214            tasks.spawn(refresh_profile(request, client));
215        }
216
217        while let Some(result) = tasks.join_next().await {
218            let (outcome, client) = match result {
219                Ok(output) => output,
220                Err(error) => {
221                    tracing::warn!(%error, "quota refresh task failed");
222                    continue;
223                }
224            };
225            if let Some(client) = client {
226                self.codex_clients
227                    .insert(outcome.report.profile_id.clone(), client);
228            }
229            self.reports
230                .insert(outcome.report.profile_id.clone(), outcome.report.clone());
231            on_report(outcome).await;
232        }
233        self.stop_clients_outside_batch(&batch).await;
234    }
235
236    /// Stop the cached clients whose profiles are not in `keep`. Every batch
237    /// carries the whole configured set, so a client left over from an earlier
238    /// batch belongs to a profile the configuration no longer has. Each one
239    /// owns a live `codex app-server` child that nothing would ever hand back
240    /// to a refresh again, so it would run until the controller exits.
241    async fn stop_clients_outside_batch(&mut self, keep: &BTreeSet<String>) {
242        let stranded = self
243            .codex_clients
244            .keys()
245            .filter(|profile_id| !keep.contains(*profile_id))
246            .cloned()
247            .collect::<Vec<_>>();
248        for profile_id in stranded {
249            if let Some(client) = self.codex_clients.remove(&profile_id) {
250                tracing::info!(profile_id, "stopping the quota client of a removed profile");
251                client.shutdown().await;
252            }
253        }
254    }
255
256    pub async fn shutdown(mut self) {
257        for (_, client) in self.codex_clients.drain() {
258            client.shutdown().await;
259        }
260    }
261}
262
263async fn refresh_profile(
264    request: QuotaRefreshRequest,
265    mut codex_client: Option<CodexUsageClient>,
266) -> (QuotaRefreshOutcome, Option<CodexUsageClient>) {
267    let credential_path = harness_authentication_marker(request.harness, &request.source_home);
268    let credential_before = credential_marker_fingerprint(&credential_path).await;
269    let QuotaRefreshRequest {
270        profile_id,
271        harness,
272        source_home,
273        environment,
274        cwd,
275    } = request;
276    let environment = environment.into_iter().collect::<HashMap<_, _>>();
277    let refreshed_at_epoch_seconds = SystemTime::now()
278        .duration_since(SystemTime::UNIX_EPOCH)
279        .unwrap_or_default()
280        .as_secs();
281    let result = match harness {
282        HarnessKind::Codex => {
283            if codex_login_is_near_expiry(&credential_path).await {
284                match codex_usage::refresh_login(
285                    &mut codex_client,
286                    cwd.clone(),
287                    environment.clone(),
288                )
289                .await
290                {
291                    Ok(()) => tracing::info!(
292                        profile_id = %profile_id,
293                        "refreshed Codex login ahead of expiry"
294                    ),
295                    Err(error) => tracing::warn!(
296                        profile_id = %profile_id,
297                        %error,
298                        "could not refresh the Codex login ahead of expiry"
299                    ),
300                }
301            }
302            let status = codex_usage::refresh(&mut codex_client, cwd, environment).await;
303            match status {
304                CodexUsageStatus::Available(report) => Ok(ProfileQuota {
305                    profile_id: profile_id.clone(),
306                    harness,
307                    windows: [report.primary, report.secondary]
308                        .into_iter()
309                        .flatten()
310                        .map(|window| QuotaWindow {
311                            label: window.label,
312                            remaining_percent: Some(window.remaining_percent),
313                            used: None,
314                            limit: None,
315                            resets: window.resets_at.and_then(format_reset_local_seconds),
316                            resets_at_epoch_seconds: window.resets_at,
317                        })
318                        .collect(),
319                    extra: None,
320                    error: None,
321                    refreshed_at_epoch_seconds,
322                }),
323                CodexUsageStatus::Unavailable(error) => Err(anyhow::anyhow!(error)),
324            }
325        }
326        HarnessKind::Claude => claude_usage::query(source_home, environment)
327            .await
328            .map(|report| ProfileQuota {
329                profile_id: profile_id.clone(),
330                harness,
331                windows: [
332                    report.five_hour.map(|window| ("5H", window)),
333                    report.week.map(|window| ("Week", window)),
334                ]
335                .into_iter()
336                .flatten()
337                .map(|(label, window)| QuotaWindow {
338                    label: label.to_string(),
339                    remaining_percent: Some(window.remaining_percent),
340                    used: None,
341                    limit: None,
342                    resets: window
343                        .reset_context
344                        .as_deref()
345                        .and_then(normalize_reset_text),
346                    resets_at_epoch_seconds: window
347                        .reset_context
348                        .as_deref()
349                        .and_then(normalize_reset_epoch_seconds),
350                })
351                .collect(),
352                extra: None,
353                error: None,
354                refreshed_at_epoch_seconds,
355            })
356            .map_err(|error| anyhow::anyhow!(error.to_string())),
357        HarnessKind::Kimi => {
358            query_kimi(&source_home, &environment)
359                .await
360                .map(|(windows, extra)| ProfileQuota {
361                    profile_id: profile_id.clone(),
362                    harness,
363                    windows,
364                    extra,
365                    error: None,
366                    refreshed_at_epoch_seconds,
367                })
368        }
369        // Grok Build publishes no HTTP quota endpoint. Its own usage view polls
370        // an ACP billing extension, and so does Mjolnir.
371        HarnessKind::Grok => {
372            grok_usage::query(source_home.clone(), cwd, environment)
373                .await
374                .map(|report| ProfileQuota {
375                    profile_id: profile_id.clone(),
376                    harness,
377                    windows: vec![QuotaWindow {
378                        label: report.period_label.clone(),
379                        remaining_percent: Some(report.remaining_percent()),
380                        // Grok Build reports a share of the allowance, not the
381                        // credit amounts behind it.
382                        used: None,
383                        limit: None,
384                        resets: report.resets_at.and_then(format_reset_local_seconds),
385                        resets_at_epoch_seconds: report.resets_at,
386                    }],
387                    extra: None,
388                    error: None,
389                    refreshed_at_epoch_seconds,
390                })
391                .map_err(|error| anyhow::anyhow!(error.to_string()))
392        }
393        HarnessKind::Deepseek => Ok(ProfileQuota {
394            profile_id: profile_id.clone(),
395            harness,
396            windows: Vec::new(),
397            extra: Some(API_LABEL.to_owned()),
398            error: None,
399            refreshed_at_epoch_seconds,
400        }),
401        HarnessKind::Muse => Ok(ProfileQuota {
402            profile_id: profile_id.clone(),
403            harness,
404            windows: Vec::new(),
405            extra: Some("Quota unavailable".into()),
406            error: None,
407            refreshed_at_epoch_seconds,
408        }),
409    };
410    let report = result.unwrap_or_else(|error| ProfileQuota {
411        profile_id,
412        harness,
413        windows: Vec::new(),
414        extra: None,
415        error: Some(error.to_string()),
416        refreshed_at_epoch_seconds,
417    });
418    let credential_after = credential_marker_fingerprint(&credential_path).await;
419    let credentials_changed = match (credential_before, credential_after) {
420        (Ok(before), Ok(after)) => before != after,
421        (Err(error), _) | (_, Err(error)) => {
422            tracing::warn!(path = %credential_path.display(), %error, "could not fingerprint quota credentials");
423            false
424        }
425    };
426    (
427        QuotaRefreshOutcome {
428            report,
429            credentials_changed,
430        },
431        codex_client,
432    )
433}
434
435/// Shortest gap to expiry Hel will leave a Codex login sitting at. A token with
436/// a long life gets a proportionally wider margin, because the poll interval
437/// buys nothing once the whole life is short.
438const CODEX_MINIMUM_REFRESH_MARGIN_MS: i64 = 60 * 60 * 1000;
439
440/// Whether the profile's Codex login is close enough to expiry that a container
441/// copy of it could reach the single-use refresh race before the next poll.
442async fn codex_login_is_near_expiry(marker: &Path) -> bool {
443    let Ok(bytes) = tokio::fs::read(marker).await else {
444        return false;
445    };
446    if bytes.len() > MAX_CREDENTIAL_BYTES {
447        return false;
448    }
449    let now = SystemTime::now()
450        .duration_since(SystemTime::UNIX_EPOCH)
451        .unwrap_or_default()
452        .as_millis() as i64;
453    codex_login_needs_refresh(
454        credential_expiry(HarnessKind::Codex, &bytes),
455        credential_freshness(HarnessKind::Codex, &bytes),
456        now,
457    )
458}
459
460/// The margin is the larger of one hour and a tenth of the token's life, where
461/// the life is what the last refresh bought. A credential that says nothing
462/// about its own age falls back to the flat hour.
463fn codex_login_needs_refresh(
464    expiry_millis: Option<i64>,
465    last_refresh_millis: Option<i64>,
466    now_millis: i64,
467) -> bool {
468    let Some(expiry) = expiry_millis else {
469        return false;
470    };
471    let lifetime = last_refresh_millis
472        .map(|refreshed| expiry.saturating_sub(refreshed))
473        .unwrap_or_default();
474    let margin = CODEX_MINIMUM_REFRESH_MARGIN_MS.max(lifetime / 10);
475    expiry.saturating_sub(now_millis) < margin
476}
477
478async fn credential_marker_fingerprint(path: &Path) -> Result<Option<String>> {
479    let metadata = match tokio::fs::metadata(path).await {
480        Ok(metadata) => metadata,
481        Err(error) if error.kind() == std::io::ErrorKind::NotFound => return Ok(None),
482        Err(error) => return Err(error).context("inspect credential marker"),
483    };
484    if metadata.len() > MAX_CREDENTIAL_BYTES as u64 {
485        bail!("credential marker exceeds {MAX_CREDENTIAL_BYTES} bytes");
486    }
487    let bytes = tokio::fs::read(path)
488        .await
489        .context("read credential marker")?;
490    if bytes.len() > MAX_CREDENTIAL_BYTES {
491        bail!("credential marker exceeds {MAX_CREDENTIAL_BYTES} bytes");
492    }
493    Ok(Some(credential_fingerprint(&bytes)))
494}
495
496async fn query_kimi(
497    home: &Path,
498    environment: &HashMap<String, String>,
499) -> Result<(Vec<QuotaWindow>, Option<String>)> {
500    let base = environment
501        .get("KIMI_CODE_BASE_URL")
502        .map(String::as_str)
503        .unwrap_or("https://api.kimi.com/coding/v1")
504        .trim_end_matches('/');
505    let client = reqwest::Client::builder()
506        .timeout(Duration::from_secs(10))
507        .build()
508        .context("build Kimi quota client")?;
509    let credentials_path = home.join("credentials/kimi-code.json");
510    let usage_url = format!("{base}/usages");
511    let response = fetch_bearer_with_auth_retry(&client, &usage_url, |force, rejected_token| {
512        ensure_fresh_kimi_token(
513            &client,
514            home,
515            &credentials_path,
516            environment,
517            force,
518            rejected_token,
519        )
520    })
521    .await?;
522    if !response.status().is_success() {
523        bail!("Kimi Code quota returned HTTP {}", response.status());
524    }
525    let payload: Value = response.json().await.context("decode Kimi Code quota")?;
526    Ok(parse_kimi_usage(&payload))
527}
528
529const KIMI_OAUTH_CLIENT_ID: &str = "17e5f671-d194-4dfb-9706-5516cb48c098";
530
531#[derive(Clone, Debug, Deserialize, Serialize)]
532struct KimiCredentials {
533    #[serde(alias = "accessToken")]
534    access_token: String,
535    #[serde(default, alias = "refreshToken")]
536    refresh_token: String,
537    #[serde(default, alias = "expiresAt")]
538    expires_at: i64,
539    #[serde(default)]
540    scope: String,
541    #[serde(default, alias = "tokenType")]
542    token_type: String,
543    #[serde(default, alias = "expiresIn")]
544    expires_in: i64,
545}
546
547impl KimiCredentials {
548    /// Whether this is a different pair from `other`. A refresh rotates the
549    /// access token, the refresh token and the expiry together, so those three
550    /// fields are what tells two pairs apart; the rest only describes them.
551    fn differs_from(&self, other: &Self) -> bool {
552        self.access_token != other.access_token
553            || self.refresh_token != other.refresh_token
554            || self.expires_at != other.expires_at
555    }
556
557    fn needs_refresh(&self) -> bool {
558        if self.expires_at == 0 {
559            return false;
560        }
561        let now = SystemTime::now()
562            .duration_since(SystemTime::UNIX_EPOCH)
563            .unwrap_or_default()
564            .as_secs() as i64;
565        let threshold = 300.max(self.expires_in / 2);
566        self.expires_at - now < threshold
567    }
568}
569
570async fn read_kimi_credentials(path: &Path) -> Result<KimiCredentials> {
571    let bytes = tokio::fs::read(path)
572        .await
573        .context("Kimi Code credentials are unavailable")?;
574    let credentials: KimiCredentials =
575        serde_json::from_slice(&bytes).context("Kimi Code credentials are invalid")?;
576    if credentials.access_token.is_empty() {
577        bail!("Kimi Code access token is missing");
578    }
579    Ok(credentials)
580}
581
582async fn fetch_bearer_with_auth_retry<F, Fut>(
583    client: &reqwest::Client,
584    url: &str,
585    mut authenticate: F,
586) -> Result<reqwest::Response>
587where
588    F: FnMut(bool, Option<String>) -> Fut,
589    Fut: std::future::Future<Output = Result<String>>,
590{
591    let token = authenticate(false, None).await?;
592    let response = client
593        .get(url)
594        .bearer_auth(&token)
595        .header(reqwest::header::ACCEPT, "application/json")
596        .send()
597        .await
598        .context("query quota")?;
599    if response.status() != reqwest::StatusCode::UNAUTHORIZED {
600        return Ok(response);
601    }
602
603    let refreshed = authenticate(true, Some(token)).await?;
604    client
605        .get(url)
606        .bearer_auth(refreshed)
607        .header(reqwest::header::ACCEPT, "application/json")
608        .send()
609        .await
610        .context("retry quota after authentication refresh")
611}
612
613/// Hand back a usable Kimi Code access token, refreshing the stored pair when
614/// it is stale or when the server rejected it.
615///
616/// The refresh lock is taken before the round trip, but a peer may break a lock
617/// it judges stale while that round trip is in flight, so ownership is checked
618/// again immediately before the credentials are written rather than only when
619/// the lock is released: see `decide_kimi_refresh_persist` for what a refresh
620/// that lost its lock does with the pair it fetched.
621async fn ensure_fresh_kimi_token(
622    client: &reqwest::Client,
623    home: &Path,
624    credentials_path: &Path,
625    environment: &HashMap<String, String>,
626    force: bool,
627    rejected_token: Option<String>,
628) -> Result<String> {
629    let initial = read_kimi_credentials(credentials_path).await?;
630    if !force && !initial.needs_refresh() {
631        return Ok(initial.access_token);
632    }
633
634    let refresh_lock = KimiRefreshLock::acquire(home).await?;
635    let active = read_kimi_credentials(credentials_path).await?;
636    let changed_while_waiting = active.differs_from(&initial);
637    if (!force && !active.needs_refresh())
638        || (force
639            && (changed_while_waiting
640                || rejected_token.is_some_and(|token| token != active.access_token)))
641    {
642        refresh_lock.release().await?;
643        return Ok(active.access_token);
644    }
645    if active.refresh_token.is_empty() {
646        refresh_lock.release().await?;
647        bail!("Kimi Code refresh token is missing; run `kimi login`");
648    }
649
650    let oauth_host = environment
651        .get("KIMI_CODE_OAUTH_HOST")
652        .or_else(|| environment.get("KIMI_OAUTH_HOST"))
653        .map(String::as_str)
654        .unwrap_or("https://auth.kimi.com")
655        .trim_end_matches('/');
656    let response = client
657        .post(format!("{oauth_host}/api/oauth/token"))
658        .header(reqwest::header::ACCEPT, "application/json")
659        .form(&[
660            ("client_id", KIMI_OAUTH_CLIENT_ID),
661            ("grant_type", "refresh_token"),
662            ("refresh_token", active.refresh_token.as_str()),
663        ])
664        .send()
665        .await
666        .context("refresh Kimi Code access token")?;
667    if !response.status().is_success() {
668        let status = response.status();
669        if matches!(
670            status,
671            reqwest::StatusCode::UNAUTHORIZED | reqwest::StatusCode::FORBIDDEN
672        ) {
673            tokio::time::sleep(Duration::from_millis(100)).await;
674            let recovery = read_kimi_credentials(credentials_path).await?;
675            if recovery.refresh_token != active.refresh_token && !recovery.access_token.is_empty() {
676                refresh_lock.release().await?;
677                return Ok(recovery.access_token);
678            }
679        }
680        refresh_lock.release().await?;
681        bail!("Kimi Code token refresh returned HTTP {status}");
682    }
683
684    let payload: Value = response
685        .json()
686        .await
687        .context("decode Kimi Code token refresh")?;
688    let access_token = required_string(&payload, "access_token", "Kimi Code token refresh")?;
689    let refresh_token = required_string(&payload, "refresh_token", "Kimi Code token refresh")?;
690    let expires_in = payload
691        .get("expires_in")
692        .and_then(value_i64)
693        .filter(|value| *value > 0)
694        .context("Kimi Code token refresh is missing expires_in")?;
695    let now = SystemTime::now()
696        .duration_since(SystemTime::UNIX_EPOCH)
697        .unwrap_or_default()
698        .as_secs() as i64;
699    let refreshed = KimiCredentials {
700        access_token: access_token.to_string(),
701        refresh_token: refresh_token.to_string(),
702        expires_at: now + expires_in,
703        scope: payload
704            .get("scope")
705            .and_then(Value::as_str)
706            .unwrap_or_default()
707            .to_string(),
708        token_type: payload
709            .get("token_type")
710            .and_then(Value::as_str)
711            .unwrap_or("Bearer")
712            .to_string(),
713        expires_in,
714    };
715    // Prove the lock is still Mjolnir's before the write, not after it: a peer that
716    // broke the lock during the round trip may already have stored a newer pair,
717    // and overwriting that would strand both refreshes.
718    let ownership = confirm_kimi_lock_ownership(&refresh_lock.path, &refresh_lock.ownership);
719    let on_disk = match &ownership {
720        // A proven loss makes the file the authority on which pair is live.
721        Err(KimiLockLoss::Stolen { .. } | KimiLockLoss::Gone) => {
722            read_kimi_credentials(credentials_path).await.ok()
723        }
724        Ok(_) | Err(KimiLockLoss::Unproven(_)) => None,
725    };
726    match decide_kimi_refresh_persist(&ownership, on_disk.as_ref(), &active) {
727        KimiRefreshPersist::Save => {
728            save_kimi_credentials(credentials_path, &refreshed)?;
729            if let Err(error) = refresh_lock.release().await {
730                // The lock was Mjolnir's when the pair was written and the write
731                // landed; losing it in the microseconds since costs the lock,
732                // not a valid credential.
733                tracing::warn!(
734                    %error,
735                    "saved refreshed Kimi Code credentials, then lost the OAuth refresh lock before releasing it"
736                );
737            }
738            Ok(refreshed.access_token)
739        }
740        KimiRefreshPersist::SaveContested(loss) => {
741            save_kimi_credentials(credentials_path, &refreshed)?;
742            tracing::warn!(
743                path = %refresh_lock.path.display(),
744                %loss,
745                "another Kimi Code token refresh took the OAuth refresh lock, but left behind the pair this refresh already spent; saved Mjolnir's refreshed pair, the only live one"
746            );
747            Ok(refreshed.access_token)
748        }
749        KimiRefreshPersist::Adopt { access_token, loss } => {
750            tracing::warn!(
751                path = %refresh_lock.path.display(),
752                %loss,
753                "another Kimi Code token refresh took the OAuth refresh lock and stored its own credentials; using those instead of the pair Mjolnir just fetched"
754            );
755            Ok(access_token)
756        }
757    }
758}
759
760fn save_kimi_credentials(path: &Path, credentials: &KimiCredentials) -> Result<()> {
761    let mut body = serde_json::to_vec_pretty(credentials)?;
762    body.push(b'\n');
763    hel::hel_config::atomic_write(path, &body).context("save refreshed Kimi Code credentials")
764}
765
766/// What a completed refresh does with the pair it just fetched.
767#[derive(Debug, Clone, PartialEq, Eq)]
768enum KimiRefreshPersist {
769    /// The lock is Mjolnir's: save the refreshed pair and give the lock back.
770    Save,
771    /// The lock is another refresher's, and the file still holds the pair this
772    /// refresh spent: save the refreshed pair anyway and leave the lock alone.
773    SaveContested(KimiLockLoss),
774    /// The lock's new holder finished first and stored its own pair: return
775    /// that token and write nothing.
776    Adopt {
777        access_token: String,
778        loss: KimiLockLoss,
779    },
780}
781
782/// Decide how a completed refresh persists its result.
783///
784/// `ownership` is the lock check taken immediately before the write, `on_disk`
785/// the pair the credentials file carried when that check reported a proven
786/// loss (`None` when the file was not consulted, or could not be read), and
787/// `active` the pair whose refresh token this refresh spent at the server.
788///
789/// A lost lock never fails the refresh: exactly one of the two pairs is live,
790/// and the file says which. A pair on disk that moved on from `active` is the
791/// other refresher's, and it is the live one, because the server rotated Mjolnir's
792/// pair away from it. A file that still holds `active` is dead whichever
793/// refresh wrote it โ€” its refresh token is the one Mjolnir just spent โ€” so the
794/// refreshed pair is the only live credential anywhere and has to be stored,
795/// even over a contested lock; leaving the spent pair in place would force a
796/// `kimi login`. The peer recovers the same way Mjolnir does, by re-reading the
797/// file when the server rejects its consumed token.
798fn decide_kimi_refresh_persist(
799    ownership: &Result<SystemTime, KimiLockLoss>,
800    on_disk: Option<&KimiCredentials>,
801    active: &KimiCredentials,
802) -> KimiRefreshPersist {
803    let loss = match ownership {
804        Ok(_) => return KimiRefreshPersist::Save,
805        // A check that could not read the directory proves nothing about who
806        // holds it, so it is no reason to treat the lock as lost.
807        Err(KimiLockLoss::Unproven(_)) => return KimiRefreshPersist::Save,
808        Err(loss) => loss.clone(),
809    };
810    match on_disk {
811        Some(pair) if pair.differs_from(active) => KimiRefreshPersist::Adopt {
812            access_token: pair.access_token.clone(),
813            loss,
814        },
815        _ => KimiRefreshPersist::SaveContested(loss),
816    }
817}
818
819fn required_string<'a>(payload: &'a Value, key: &str, context: &str) -> Result<&'a str> {
820    payload
821        .get(key)
822        .and_then(Value::as_str)
823        .filter(|value| !value.is_empty())
824        .with_context(|| format!("{context} is missing {key}"))
825}
826
827struct KimiRefreshLock {
828    path: std::path::PathBuf,
829    ownership: Arc<Mutex<KimiLockOwnership>>,
830    heartbeat: Option<tokio::task::JoinHandle<()>>,
831}
832
833/// What Mjolnir knows about the lock directory it created. The Kimi Code CLI
834/// breaks a lock whose modification time stopped moving and takes it over, so
835/// holding the directory is not the same as owning it: Mjolnir checks the mtime it
836/// published is still there before touching or removing the directory.
837/// Touching a lock the CLI now owns trips the CLI's own ownership check
838/// (`ECOMPROMISED`, proper-lockfile 4.1.2 `lib/lockfile.js:114-140`), and
839/// removing it would hand a third holder a lock the CLI is still using.
840#[derive(Debug)]
841enum KimiLockOwnership {
842    /// Mjolnir published this modification time and the directory still carried it
843    /// when Mjolnir last looked.
844    Held(SystemTime),
845    /// Mjolnir must not touch or remove the directory again.
846    Lost(KimiLockLoss),
847}
848
849/// Why the lock directory is not Mjolnir's any more.
850#[derive(Debug, Clone, PartialEq, Eq)]
851enum KimiLockLoss {
852    /// It carries a modification time Mjolnir never published: another holder broke
853    /// the lock and took it.
854    Stolen {
855        published: SystemTime,
856        observed: SystemTime,
857    },
858    /// It is gone: another holder broke the lock, or Mjolnir already released it.
859    Gone,
860    /// It could not be inspected, so Mjolnir cannot prove the lock is still its
861    /// own. Mjolnir leaves it alone; whoever wants it next breaks it once Mjolnir's
862    /// modification time goes stale.
863    Unproven(String),
864}
865
866impl std::fmt::Display for KimiLockLoss {
867    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
868        match self {
869            Self::Stolen {
870                published,
871                observed,
872            } => write!(
873                formatter,
874                "another process took it: Mjolnir published modification time {}, but the directory carries {}",
875                epoch_label(*published),
876                epoch_label(*observed)
877            ),
878            Self::Gone => formatter.write_str("another process removed it"),
879            Self::Unproven(error) => {
880                write!(
881                    formatter,
882                    "Mjolnir could not confirm it still owns it: {error}"
883                )
884            }
885        }
886    }
887}
888
889fn epoch_label(time: SystemTime) -> String {
890    match time.duration_since(SystemTime::UNIX_EPOCH) {
891        Ok(since) => format!("{:.3}", since.as_secs_f64()),
892        Err(_) => "before the epoch".to_string(),
893    }
894}
895
896/// The Kimi Code CLI is the other holder of this lock, and it agrees that a
897/// live holder keeps the directory's mtime moving. It takes the lock through
898/// `proper-lockfile` with `stale: 5_000` (kimi-code
899/// `packages/oauth/src/oauth-manager.ts:216-220`; the shipped binary carries
900/// the same `stale: 5e3`), which rewrites the mtime every `stale / 2` for as
901/// long as the lock is held (proper-lockfile 4.1.2 `lib/lockfile.js:99-183`,
902/// interval resolved at `lib/lockfile.js:220-221`) and removes any lock whose
903/// mtime is older than `stale` (`lib/lockfile.js:67-79, 84-86`). A CLI refresh
904/// can hold the lock far longer than that โ€” three tries against a 30s HTTP
905/// timeout plus backoff (`packages/oauth/src/oauth.ts:56-73, 226-263`) โ€” but
906/// never silently, so a stopped mtime still means the holder is gone. Its
907/// mtimes can also land up to a second in the future
908/// (`lib/mtime-precision.js:44-52`), which `break_stale_kimi_lock` reads as
909/// "not stale" because `duration_since` fails: the safe answer.
910const KIMI_CLI_LOCK_STALE_AFTER: Duration = Duration::from_secs(5);
911/// A holder republishes the lock directory's modification time on this
912/// interval, so a lock whose mtime stopped moving has no live holder. The CLI
913/// judges Mjolnir's lock by that same mtime, so the interval has to fit inside
914/// `KIMI_CLI_LOCK_STALE_AFTER` several times over: one beat pays for the wait
915/// between touches, and the rest is stall the heartbeat task may absorb.
916const KIMI_LOCK_HEARTBEAT_INTERVAL: Duration = Duration::from_secs(1);
917/// How long the heartbeat task may stall โ€” descheduled, starved, or blocked on
918/// a slow filesystem โ€” before the CLI is entitled to break a lock Mjolnir still
919/// holds and rotate the credentials alongside it. It is the CLI's window less
920/// the interval Mjolnir already spends waiting between touches.
921const KIMI_LOCK_HEARTBEAT_STALL_TOLERANCE: Duration = Duration::from_secs(
922    KIMI_CLI_LOCK_STALE_AFTER.as_secs() - KIMI_LOCK_HEARTBEAT_INTERVAL.as_secs(),
923);
924/// Beats of silence Mjolnir waits out before it calls another holder's lock
925/// abandoned. Several beats of slack, so a live holder delayed by the scheduler
926/// keeps its lock, and deliberately more patient than the CLI's 5s: breaking
927/// later than the peer can never steal a live lock, and it costs no recovery
928/// time, because the CLI reclaims a lock a crashed Mjolnir left behind after its
929/// own 5s.
930const KIMI_LOCK_STALE_HEARTBEATS: u64 = 10;
931/// Derived from the heartbeat so the two cannot drift apart.
932const KIMI_LOCK_STALE_AFTER: Duration =
933    Duration::from_secs(KIMI_LOCK_STALE_HEARTBEATS * KIMI_LOCK_HEARTBEAT_INTERVAL.as_secs());
934const _: () = assert!(
935    KIMI_LOCK_HEARTBEAT_STALL_TOLERANCE.as_secs() >= 4 * KIMI_LOCK_HEARTBEAT_INTERVAL.as_secs(),
936    "Mjolnir must be able to miss several beats in a row and still hold a lock the Kimi Code CLI could otherwise break"
937);
938const _: () = assert!(
939    KIMI_LOCK_STALE_AFTER.as_secs() >= KIMI_CLI_LOCK_STALE_AFTER.as_secs(),
940    "Mjolnir must not call a lock stale sooner than the Kimi Code CLI does, or it can break a lock the CLI still holds"
941);
942const KIMI_LOCK_RETRY_INTERVAL: Duration = Duration::from_millis(500);
943const KIMI_LOCK_WAIT: Duration = Duration::from_secs(60);
944/// Filesystems record modification times at their own precision โ€” whole
945/// seconds on ext3 and HFS+ โ€” and the Kimi Code CLI leans on that, rounding
946/// its own writes up to the next whole second so a coarse filesystem stores
947/// them unchanged (`lib/mtime-precision.js:44-52`); its mtimes therefore land
948/// up to a second in the future. So a time Mjolnir published and the time it reads
949/// back can differ by anything under a second and still be the same write.
950/// Nothing smaller than a second distinguishes holders: taking the lock from
951/// Mjolnir costs another holder at least `KIMI_CLI_LOCK_STALE_AFTER` of silence
952/// first, so a thief's modification time is seconds away, never milliseconds.
953const KIMI_LOCK_MTIME_TOLERANCE: Duration = Duration::from_secs(1);
954
955impl KimiRefreshLock {
956    async fn acquire(home: &Path) -> Result<Self> {
957        Self::acquire_within(home, KIMI_LOCK_WAIT).await
958    }
959
960    async fn acquire_within(home: &Path, wait: Duration) -> Result<Self> {
961        let oauth_dir = home.join("oauth");
962        tokio::fs::create_dir_all(&oauth_dir)
963            .await
964            .context("prepare Kimi Code OAuth lock")?;
965        let sentinel = oauth_dir.join("kimi-code");
966        tokio::fs::OpenOptions::new()
967            .create(true)
968            .append(true)
969            .open(&sentinel)
970            .await
971            .context("prepare Kimi Code OAuth lock sentinel")?;
972        let path = oauth_dir.join("kimi-code.lock");
973        let deadline = tokio::time::Instant::now() + wait;
974        loop {
975            match tokio::fs::create_dir(&path).await {
976                Ok(()) => return Self::claim(path).await,
977                Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
978                    if tokio::time::Instant::now() >= deadline {
979                        break;
980                    }
981                    // A holder killed mid-refresh leaves its directory behind
982                    // forever; break the lock once its heartbeat has stopped
983                    // and retry the create immediately.
984                    if !break_stale_kimi_lock(&path).await {
985                        tokio::time::sleep(KIMI_LOCK_RETRY_INTERVAL).await;
986                    }
987                }
988                Err(error) => return Err(error).context("acquire Kimi Code OAuth refresh lock"),
989            }
990        }
991        bail!(
992            "timed out waiting for Kimi Code OAuth refresh lock {}; another Kimi Code refresh is holding it, or a crashed one left it behind and the directory has to be removed",
993            path.display()
994        )
995    }
996
997    /// Take ownership of a directory Mjolnir just created. The modification time
998    /// the filesystem recorded for the create is the first proof of ownership;
999    /// every heartbeat republishes it.
1000    async fn claim(path: std::path::PathBuf) -> Result<Self> {
1001        let published = kimi_lock_mtime(&path)
1002            .map_err(anyhow::Error::new)
1003            .and_then(|mtime| mtime.context("it vanished as Mjolnir created it"));
1004        match published {
1005            Ok(published) => Ok(Self::held(path, published)),
1006            Err(error) => {
1007                // Without a first modification time Mjolnir could never prove the
1008                // lock is its own, so it could never release it either. Give it
1009                // back now instead of leaving it for a stale-breaker.
1010                let _ = tokio::fs::remove_dir(&path).await;
1011                Err(error).with_context(|| {
1012                    format!(
1013                        "claim the new Kimi Code OAuth refresh lock {}",
1014                        path.display()
1015                    )
1016                })
1017            }
1018        }
1019    }
1020
1021    fn held(path: std::path::PathBuf, published: SystemTime) -> Self {
1022        let ownership = Arc::new(Mutex::new(KimiLockOwnership::Held(published)));
1023        let heartbeat_path = path.clone();
1024        let heartbeat_ownership = Arc::clone(&ownership);
1025        let heartbeat = tokio::spawn(async move {
1026            loop {
1027                tokio::time::sleep(KIMI_LOCK_HEARTBEAT_INTERVAL).await;
1028                match beat_kimi_lock(&heartbeat_path, &heartbeat_ownership) {
1029                    Ok(()) => {}
1030                    Err(KimiLockLoss::Unproven(error)) => {
1031                        tracing::debug!(path = %heartbeat_path.display(), %error, "heartbeat Kimi Code OAuth refresh lock");
1032                    }
1033                    Err(loss) => {
1034                        tracing::warn!(path = %heartbeat_path.display(), %loss, "stopped heartbeating a Kimi Code OAuth refresh lock Mjolnir no longer holds");
1035                        return;
1036                    }
1037                }
1038            }
1039        });
1040        Self {
1041            path,
1042            ownership,
1043            heartbeat: Some(heartbeat),
1044        }
1045    }
1046
1047    /// Give the lock back. Fails when the lock stopped being Mjolnir's, because the
1048    /// refresh it was protecting then ran beside another one. Callers that have
1049    /// already confirmed ownership and stored valid credentials treat that
1050    /// failure as a lost lock rather than a failed refresh; see
1051    /// `ensure_fresh_kimi_token`.
1052    async fn release(mut self) -> Result<()> {
1053        if let Some(heartbeat) = self.heartbeat.take() {
1054            heartbeat.abort();
1055        }
1056        if let Err(loss) = confirm_kimi_lock_ownership(&self.path, &self.ownership) {
1057            bail!(
1058                "the Kimi Code OAuth refresh lock {} stopped being Mjolnir's mid-refresh: {loss}; another Kimi Code token refresh may have rotated the credentials beside this one",
1059                self.path.display()
1060            );
1061        }
1062        match tokio::fs::remove_dir(&self.path).await {
1063            Ok(()) => {
1064                *lock_ownership(&self.ownership) = KimiLockOwnership::Lost(KimiLockLoss::Gone)
1065            }
1066            Err(error) => {
1067                tracing::warn!(path = %self.path.display(), %error, "release Kimi Code OAuth refresh lock");
1068            }
1069        }
1070        Ok(())
1071    }
1072}
1073
1074impl Drop for KimiRefreshLock {
1075    fn drop(&mut self) {
1076        if let Some(heartbeat) = self.heartbeat.take() {
1077            heartbeat.abort();
1078        }
1079        if matches!(*lock_ownership(&self.ownership), KimiLockOwnership::Lost(_)) {
1080            // Already released, or reported where the loss was discovered.
1081            return;
1082        }
1083        match confirm_kimi_lock_ownership(&self.path, &self.ownership) {
1084            Ok(_) => {
1085                if let Err(error) = std::fs::remove_dir(&self.path) {
1086                    tracing::warn!(path = %self.path.display(), %error, "release Kimi Code OAuth refresh lock");
1087                }
1088            }
1089            Err(loss) => {
1090                tracing::warn!(path = %self.path.display(), %loss, "left a Kimi Code OAuth refresh lock Mjolnir no longer holds in place");
1091            }
1092        }
1093    }
1094}
1095
1096fn lock_ownership(
1097    ownership: &Mutex<KimiLockOwnership>,
1098) -> std::sync::MutexGuard<'_, KimiLockOwnership> {
1099    ownership
1100        .lock()
1101        .unwrap_or_else(|poisoned| poisoned.into_inner())
1102}
1103
1104/// Check that the lock directory still carries the modification time Mjolnir
1105/// published, and remember a loss so every later check agrees. `Ok` hands back
1106/// the published time; `Err` means Mjolnir must neither touch nor remove the
1107/// directory.
1108fn confirm_kimi_lock_ownership(
1109    path: &Path,
1110    ownership: &Mutex<KimiLockOwnership>,
1111) -> Result<SystemTime, KimiLockLoss> {
1112    let published = match &*lock_ownership(ownership) {
1113        KimiLockOwnership::Held(published) => *published,
1114        KimiLockOwnership::Lost(loss) => return Err(loss.clone()),
1115    };
1116    let loss = match kimi_lock_mtime(path) {
1117        Ok(Some(observed)) if kimi_lock_mtime_matches(published, observed) => {
1118            return Ok(published);
1119        }
1120        Ok(Some(observed)) => KimiLockLoss::Stolen {
1121            published,
1122            observed,
1123        },
1124        Ok(None) => KimiLockLoss::Gone,
1125        // A stat that fails says nothing about who holds the lock, so the
1126        // ownership Mjolnir recorded stands and a later beat can confirm it again.
1127        Err(error) => return Err(KimiLockLoss::Unproven(error.to_string())),
1128    };
1129    *lock_ownership(ownership) = KimiLockOwnership::Lost(loss.clone());
1130    Err(loss)
1131}
1132
1133/// One heartbeat: prove the directory is still Mjolnir's, then publish a fresh
1134/// modification time on it.
1135fn beat_kimi_lock(path: &Path, ownership: &Mutex<KimiLockOwnership>) -> Result<(), KimiLockLoss> {
1136    confirm_kimi_lock_ownership(path, ownership)?;
1137    let published = SystemTime::now();
1138    if let Err(error) = touch_kimi_lock(path, published) {
1139        // Only a time actually written may be remembered, or the next check
1140        // would report a theft that never happened.
1141        return Err(KimiLockLoss::Unproven(error.to_string()));
1142    }
1143    let mut ownership = lock_ownership(ownership);
1144    if matches!(*ownership, KimiLockOwnership::Held(_)) {
1145        *ownership = KimiLockOwnership::Held(published);
1146    }
1147    Ok(())
1148}
1149
1150/// Whether a modification time read back from the lock directory is the one Mjolnir
1151/// published. See `KIMI_LOCK_MTIME_TOLERANCE` for why a sub-second difference
1152/// is the same write rather than another holder's.
1153fn kimi_lock_mtime_matches(published: SystemTime, observed: SystemTime) -> bool {
1154    observed
1155        .duration_since(published)
1156        .or_else(|_| published.duration_since(observed))
1157        .is_ok_and(|drift| drift < KIMI_LOCK_MTIME_TOLERANCE)
1158}
1159
1160/// The lock directory's modification time, or `None` when the directory is
1161/// gone. Mjolnir and the Kimi Code CLI share this one value and nothing else: the
1162/// CLI releases its lock with a plain `rmdir`, so the directory has to stay
1163/// empty and the mtime is the whole protocol.
1164fn kimi_lock_mtime(path: &Path) -> std::io::Result<Option<SystemTime>> {
1165    match std::fs::metadata(path) {
1166        Ok(metadata) => metadata.modified().map(Some),
1167        Err(error) if error.kind() == std::io::ErrorKind::NotFound => Ok(None),
1168        Err(error) => Err(error),
1169    }
1170}
1171
1172/// Publish a lock directory's modification time, the signal that its holder is
1173/// still alive. Windows opens a directory handle only under backup semantics,
1174/// so the heartbeat would otherwise be a silent no-op there and every live lock
1175/// would look abandoned.
1176fn touch_kimi_lock(path: &Path, modified: SystemTime) -> std::io::Result<()> {
1177    let mut options = std::fs::OpenOptions::new();
1178    options.read(true);
1179    #[cfg(windows)]
1180    {
1181        use std::os::windows::fs::OpenOptionsExt;
1182        const FILE_FLAG_BACKUP_SEMANTICS: u32 = 0x0200_0000;
1183        options.custom_flags(FILE_FLAG_BACKUP_SEMANTICS);
1184    }
1185    options
1186        .open(path)?
1187        .set_times(std::fs::FileTimes::new().set_modified(modified))
1188}
1189
1190/// Remove a lock directory whose heartbeat has stopped, so a holder killed
1191/// mid-refresh cannot poison the profile home until someone removes it by hand.
1192/// Returns whether the lock is gone and the caller should retry the create at
1193/// once; a lock another process removes or recreates underneath simply loses or
1194/// wins the next create.
1195async fn break_stale_kimi_lock(path: &Path) -> bool {
1196    let modified = match kimi_lock_mtime(path) {
1197        Ok(Some(modified)) => modified,
1198        Ok(None) => return true,
1199        Err(error) => {
1200            tracing::warn!(path = %path.display(), %error, "inspect Kimi Code OAuth refresh lock");
1201            return false;
1202        }
1203    };
1204    let age = SystemTime::now().duration_since(modified).ok();
1205    let Some(age) = age.filter(|age| *age >= KIMI_LOCK_STALE_AFTER) else {
1206        return false;
1207    };
1208    match tokio::fs::remove_dir(path).await {
1209        Ok(()) => {
1210            tracing::warn!(
1211                path = %path.display(),
1212                age_seconds = age.as_secs(),
1213                "removed a Kimi Code OAuth refresh lock whose holder stopped heartbeating"
1214            );
1215            true
1216        }
1217        Err(error) if error.kind() == std::io::ErrorKind::NotFound => true,
1218        Err(error) => {
1219            tracing::warn!(path = %path.display(), %error, "remove stale Kimi Code OAuth refresh lock");
1220            false
1221        }
1222    }
1223}
1224
1225fn parse_kimi_usage(payload: &Value) -> (Vec<QuotaWindow>, Option<String>) {
1226    let mut windows = Vec::new();
1227    if let Some(summary) = payload.get("usage")
1228        && let Some(window) = parse_kimi_window(summary, "Weekly limit")
1229    {
1230        windows.push(window);
1231    }
1232    if let Some(limits) = payload.get("limits").and_then(Value::as_array) {
1233        for (index, item) in limits.iter().enumerate() {
1234            let detail = item.get("detail").unwrap_or(item);
1235            if let Some(window) = parse_kimi_window(detail, &format!("Limit #{}", index + 1)) {
1236                windows.push(window);
1237            }
1238        }
1239    }
1240    let extra = payload
1241        .pointer("/boosterWallet/balance/amountLeft")
1242        .and_then(value_i64)
1243        .map(|value| format!("booster {} remaining", value / 1_000_000));
1244    (windows, extra)
1245}
1246
1247fn parse_kimi_window(value: &Value, fallback: &str) -> Option<QuotaWindow> {
1248    let limit = value.get("limit").and_then(value_i64);
1249    let used = value.get("used").and_then(value_i64).or_else(|| {
1250        let remaining = value.get("remaining").and_then(value_i64)?;
1251        Some(limit? - remaining)
1252    });
1253    if used.is_none() && limit.is_none() {
1254        return None;
1255    }
1256    let provider_label = value
1257        .get("name")
1258        .or_else(|| value.get("title"))
1259        .and_then(Value::as_str)
1260        .unwrap_or(fallback);
1261    let label = if provider_label.to_ascii_lowercase().contains("week") {
1262        "Week".to_string()
1263    } else if provider_label.to_ascii_lowercase().contains("5h") || fallback.starts_with("Limit #")
1264    {
1265        "5H".to_string()
1266    } else {
1267        provider_label.to_string()
1268    };
1269    let reset_value = ["resetAt", "reset_at", "resetTime", "reset_time"]
1270        .iter()
1271        .find_map(|key| value.get(*key));
1272    let resets = reset_value.and_then(normalize_kimi_reset);
1273    let resets_at_epoch_seconds = reset_value.and_then(kimi_reset_epoch_seconds);
1274    let remaining_percent = match (used, limit) {
1275        (Some(used), Some(limit)) if limit > 0 => {
1276            Some((100 - used.clamp(0, limit) * 100 / limit) as u8)
1277        }
1278        _ => None,
1279    };
1280    Some(QuotaWindow {
1281        label,
1282        remaining_percent,
1283        used,
1284        limit,
1285        resets,
1286        resets_at_epoch_seconds,
1287    })
1288}
1289
1290fn value_i64(value: &Value) -> Option<i64> {
1291    value
1292        .as_i64()
1293        .or_else(|| value.as_str()?.parse::<i64>().ok())
1294}
1295
1296fn normalize_kimi_reset(value: &Value) -> Option<String> {
1297    value
1298        .as_f64()
1299        .and_then(format_reset_local)
1300        .or_else(|| value.as_str().and_then(normalize_reset_text))
1301}
1302
1303fn kimi_reset_epoch_seconds(value: &Value) -> Option<i64> {
1304    value
1305        .as_f64()
1306        .map(|epoch| {
1307            if epoch.abs() >= 1_000_000_000_000.0 {
1308                (epoch / 1000.0).trunc() as i64
1309            } else {
1310                epoch.trunc() as i64
1311            }
1312        })
1313        .or_else(|| value.as_str().and_then(normalize_reset_epoch_seconds))
1314}
1315
1316/// Format a Unix reset timestamp as wall-clock time in the machine's local
1317/// time zone. Accepts seconds or milliseconds and rejects non-finite or
1318/// out-of-range values.
1319pub(crate) fn format_reset_local(epoch: f64) -> Option<String> {
1320    if !epoch.is_finite() {
1321        return None;
1322    }
1323    let seconds = if epoch.abs() >= 1_000_000_000_000.0 {
1324        (epoch / 1000.0).trunc() as i64
1325    } else {
1326        epoch.trunc() as i64
1327    };
1328    let local = Local.timestamp_opt(seconds, 0).single()?;
1329    Some(format_reset_label(local.fixed_offset()))
1330}
1331
1332pub(crate) fn format_reset_local_seconds(epoch: i64) -> Option<String> {
1333    format_reset_local(epoch as f64)
1334}
1335
1336/// Normalize a provider's textual reset value to the compact 24-hour form
1337/// used by the dashboard. A time-only value is the next occurrence of that
1338/// wall-clock time; Claude Code uses this shape for its five-hour window.
1339pub(crate) fn normalize_reset_text(value: &str) -> Option<String> {
1340    normalize_reset_at(value, Local::now().fixed_offset()).map(format_reset_label)
1341}
1342
1343pub(crate) fn normalize_reset_epoch_seconds(value: &str) -> Option<i64> {
1344    normalize_reset_at(value, Local::now().fixed_offset()).map(|reset| reset.timestamp())
1345}
1346
1347fn normalize_reset_at(value: &str, now: DateTime<FixedOffset>) -> Option<DateTime<FixedOffset>> {
1348    let value = value.trim();
1349    if value.is_empty() {
1350        return None;
1351    }
1352    if let Ok(epoch) = value.parse::<f64>() {
1353        let seconds = if epoch.abs() >= 1_000_000_000_000.0 {
1354            (epoch / 1000.0).trunc() as i64
1355        } else {
1356            epoch.trunc() as i64
1357        };
1358        return Local
1359            .timestamp_opt(seconds, 0)
1360            .single()
1361            .map(|reset| reset.fixed_offset());
1362    }
1363    if let Ok(timestamp) = DateTime::parse_from_rfc3339(value) {
1364        return Some(timestamp.with_timezone(&Local).fixed_offset());
1365    }
1366
1367    let value = value
1368        .strip_prefix("at ")
1369        .unwrap_or(value)
1370        .split('(')
1371        .next()
1372        .unwrap_or(value)
1373        .trim()
1374        .trim_end_matches(',');
1375    let parse_time = |value: &str| {
1376        let value = value
1377            .to_ascii_lowercase()
1378            .chars()
1379            .filter(|ch| !ch.is_whitespace())
1380            .collect::<String>();
1381        let value = ["am", "pm"]
1382            .into_iter()
1383            .find_map(|suffix| {
1384                let hour = value.strip_suffix(suffix)?;
1385                (!hour.contains(':')).then(|| format!("{hour}:00{suffix}"))
1386            })
1387            .unwrap_or(value);
1388        ["%I:%M%P", "%I%P", "%H:%M"]
1389            .iter()
1390            .find_map(|format| NaiveTime::parse_from_str(&value, format).ok())
1391    };
1392
1393    // Claude has used both `Aug 14 at 4am` and `Aug 14, 4am` across
1394    // releases. Keep the provider punctuation out of the date/time parsers.
1395    let dated_time = value.split_once(" at ").or_else(|| {
1396        value
1397            .split_once(',')
1398            .map(|(date, time)| (date, time.trim()))
1399    });
1400    if let Some((date, time)) = dated_time {
1401        let time = parse_time(time.trim())?;
1402        let date = date.trim().trim_end_matches(',');
1403        let date = match date.to_ascii_lowercase().as_str() {
1404            "today" => now.date_naive(),
1405            "tomorrow" => now.date_naive().checked_add_days(Days::new(1))?,
1406            _ => NaiveDate::parse_from_str(
1407                &format!("{} {}", date.replace(',', ""), now.year()),
1408                "%b %e %Y",
1409            )
1410            .ok()?,
1411        };
1412        return now
1413            .timezone()
1414            .from_local_datetime(&date.and_time(time))
1415            .single();
1416    }
1417
1418    let time = parse_time(value)?;
1419    let mut date = now.date_naive();
1420    let mut reset = now
1421        .timezone()
1422        .from_local_datetime(&date.and_time(time))
1423        .single()?;
1424    if reset <= now {
1425        date = date.checked_add_days(Days::new(1))?;
1426        reset = now
1427            .timezone()
1428            .from_local_datetime(&date.and_time(time))
1429            .single()?;
1430    }
1431    Some(reset)
1432}
1433
1434/// Pure formatter split from local-zone discovery for deterministic tests.
1435fn format_reset_label(reset: DateTime<FixedOffset>) -> String {
1436    reset.format("%H:%M %b %-d").to_string()
1437}
1438
1439#[cfg(test)]
1440mod tests {
1441    use super::*;
1442    use axum::body::Bytes;
1443    use axum::extract::State;
1444    use axum::http::{HeaderMap, StatusCode};
1445    use axum::routing::{get, post};
1446    use axum::{Json, Router};
1447    use std::sync::{Arc, Mutex};
1448
1449    #[test]
1450    fn parses_kimi_summary_limits_and_booster_without_credentials() {
1451        let payload = serde_json::json!({
1452            "usage": {"name":"Weekly", "used":40, "limit":1000, "resetAt":"tomorrow"},
1453            "limits": [{"detail":{"remaining":"90", "limit":"100", "name":"5h"}}],
1454            "boosterWallet": {"balance":{"amountLeft":42000000}}
1455        });
1456        let (windows, extra) = parse_kimi_usage(&payload);
1457        assert_eq!(windows.len(), 2);
1458        assert_eq!(windows[0].used, Some(40));
1459        assert_eq!(windows[1].used, Some(10));
1460        assert_eq!(windows[0].label, "Week");
1461        assert_eq!(windows[0].remaining_percent, Some(96));
1462        assert_eq!(windows[1].label, "5H");
1463        assert_eq!(windows[1].remaining_percent, Some(90));
1464        assert_eq!(extra.as_deref(), Some("booster 42 remaining"));
1465    }
1466
1467    #[test]
1468    fn compact_includes_reset_and_error_states() {
1469        let report = ProfileQuota {
1470            profile_id: "codex-1".into(),
1471            harness: HarnessKind::Codex,
1472            windows: vec![QuotaWindow {
1473                label: "5H".into(),
1474                remaining_percent: Some(70),
1475                used: None,
1476                limit: None,
1477                resets: Some("10:00 Jun 17".into()),
1478                resets_at_epoch_seconds: Some(14_400),
1479            }],
1480            extra: None,
1481            error: None,
1482            refreshed_at_epoch_seconds: 0,
1483        };
1484        assert!(report.compact().contains("70% left"));
1485        assert!(report.compact().contains("resets 10:00 Jun 17"));
1486    }
1487
1488    #[test]
1489    fn compact_shows_login_expired_without_unavailable_prefix() {
1490        let report = ProfileQuota {
1491            profile_id: "claude2".into(),
1492            harness: HarnessKind::Claude,
1493            windows: vec![],
1494            extra: None,
1495            error: Some(claude_usage::LOGIN_EXPIRED.into()),
1496            refreshed_at_epoch_seconds: 0,
1497        };
1498        assert_eq!(report.compact(), claude_usage::LOGIN_EXPIRED);
1499        assert_eq!(
1500            report.error_label().as_deref(),
1501            Some(claude_usage::LOGIN_EXPIRED)
1502        );
1503    }
1504
1505    #[test]
1506    fn compact_still_prefixes_other_errors_with_unavailable() {
1507        let report = ProfileQuota {
1508            profile_id: "claude2".into(),
1509            harness: HarnessKind::Claude,
1510            windows: vec![],
1511            extra: None,
1512            error: Some("query Claude usage: HTTP 429".into()),
1513            refreshed_at_epoch_seconds: 0,
1514        };
1515        assert_eq!(
1516            report.compact(),
1517            "unavailable: query Claude usage: HTTP 429"
1518        );
1519    }
1520
1521    #[test]
1522    fn compact_displays_a_shared_reset_once() {
1523        let report = ProfileQuota {
1524            profile_id: "codex-1".into(),
1525            harness: HarnessKind::Codex,
1526            windows: vec![
1527                QuotaWindow {
1528                    label: "5H".into(),
1529                    remaining_percent: Some(70),
1530                    used: None,
1531                    limit: None,
1532                    resets: Some("10:00 Jun 17".into()),
1533                    resets_at_epoch_seconds: Some(14_400),
1534                },
1535                QuotaWindow {
1536                    label: "Week".into(),
1537                    remaining_percent: Some(55),
1538                    used: None,
1539                    limit: None,
1540                    resets: Some("10:00 Jun 17".into()),
1541                    resets_at_epoch_seconds: Some(14_400),
1542                },
1543            ],
1544            extra: None,
1545            error: None,
1546            refreshed_at_epoch_seconds: 0,
1547        };
1548        assert_eq!(
1549            report.compact(),
1550            "5H 70% left, resets 10:00 Jun 17 ยท Week 55% left"
1551        );
1552    }
1553
1554    #[test]
1555    fn compact_hides_claude_short_window_when_week_is_exhausted() {
1556        let report = ProfileQuota {
1557            profile_id: "claude".into(),
1558            harness: HarnessKind::Claude,
1559            windows: vec![
1560                QuotaWindow {
1561                    label: "5H".into(),
1562                    remaining_percent: Some(100),
1563                    used: None,
1564                    limit: None,
1565                    resets: None,
1566                    resets_at_epoch_seconds: None,
1567                },
1568                QuotaWindow {
1569                    label: "Week".into(),
1570                    remaining_percent: Some(0),
1571                    used: None,
1572                    limit: None,
1573                    resets: Some("03:59 Aug 14".into()),
1574                    resets_at_epoch_seconds: None,
1575                },
1576            ],
1577            extra: None,
1578            error: None,
1579            refreshed_at_epoch_seconds: 0,
1580        };
1581
1582        assert_eq!(report.compact(), "Week 0% left, resets 03:59 Aug 14");
1583    }
1584
1585    #[cfg(unix)]
1586    #[tokio::test]
1587    async fn a_grok_profile_reports_its_billing_period_as_one_quota_window() {
1588        use std::os::unix::fs::PermissionsExt;
1589
1590        let directory = tempfile::tempdir().unwrap();
1591        let executable = directory.path().join("grok");
1592        std::fs::write(directory.path().join("auth.json"), b"old credentials").unwrap();
1593        std::fs::write(
1594            &executable,
1595            "#!/bin/sh\nprintf 'refreshed credentials' > \"$GROK_HOME/auth.json\"\nwhile IFS= read -r line; do\n  case \"$line\" in\n    *initialize*) printf '{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{}}\\n' ;;\n    *billing*) printf '{\"jsonrpc\":\"2.0\",\"id\":2,\"result\":{\"config\":{\"creditUsagePercent\":25.0,\"currentPeriod\":{\"type\":\"USAGE_PERIOD_TYPE_WEEKLY\",\"end\":\"2026-08-18T05:22:07+00:00\"}},\"subscription_tier\":\"X Premium+\"}}\\n' ;;\n  esac\ndone\n",
1596        )
1597        .unwrap();
1598        std::fs::set_permissions(&executable, std::fs::Permissions::from_mode(0o755)).unwrap();
1599        let environment = BTreeMap::from([
1600            (
1601                "GROK_HOME".to_owned(),
1602                directory.path().to_string_lossy().into_owned(),
1603            ),
1604            (
1605                "PATH".to_owned(),
1606                directory.path().to_string_lossy().into_owned(),
1607            ),
1608        ]);
1609
1610        let (outcome, _) = refresh_profile(
1611            QuotaRefreshRequest {
1612                profile_id: "grok".into(),
1613                harness: HarnessKind::Grok,
1614                source_home: directory.path().to_path_buf(),
1615                environment,
1616                cwd: directory.path().to_path_buf(),
1617            },
1618            None,
1619        )
1620        .await;
1621        assert!(outcome.credentials_changed);
1622        let report = outcome.report;
1623
1624        assert_eq!(report.error, None, "{:?}", report.error);
1625        // One long window and no short one: Grok Build has no 5-hour budget.
1626        assert_eq!(report.windows.len(), 1);
1627        assert_eq!(report.weekly_window().unwrap().remaining_percent, Some(75));
1628        assert_eq!(report.five_hour_window(), None);
1629        // The subscription tier stays off the row; the fixture carries it to
1630        // prove it is ignored.
1631        assert_eq!(report.extra, None);
1632        assert!(report.compact().starts_with("Week 75% left, resets "));
1633    }
1634
1635    /// A `codex app-server` stand-in on `PATH` that logs every request line it
1636    /// reads, so a test can assert the exact protocol exchange.
1637    #[cfg(unix)]
1638    fn fake_codex_app_server(
1639        directory: &Path,
1640        script: &str,
1641    ) -> (BTreeMap<String, String>, std::path::PathBuf) {
1642        use std::os::unix::fs::PermissionsExt;
1643
1644        let executable = directory.join("codex");
1645        std::fs::write(&executable, script).unwrap();
1646        std::fs::set_permissions(&executable, std::fs::Permissions::from_mode(0o755)).unwrap();
1647        let log = directory.join("requests.jsonl");
1648        let environment = BTreeMap::from([
1649            ("PATH".to_owned(), directory.to_string_lossy().into_owned()),
1650            (
1651                "CODEX_USAGE_TEST_LOG".to_owned(),
1652                log.to_string_lossy().into_owned(),
1653            ),
1654            (
1655                "CODEX_AUTH_FILE".to_owned(),
1656                directory.join("auth.json").to_string_lossy().into_owned(),
1657            ),
1658        ]);
1659        (environment, log)
1660    }
1661
1662    /// A Codex `auth.json` whose access token is a JWT expiring `expires_in`
1663    /// from now, last refreshed `refreshed_ago` before now.
1664    #[cfg(unix)]
1665    fn write_codex_auth(home: &Path, expires_in: Duration, refreshed_ago: Duration) {
1666        use base64::Engine as _;
1667
1668        let now = chrono::Utc::now();
1669        let expiry = (now + chrono::TimeDelta::from_std(expires_in).unwrap()).timestamp();
1670        let segment = |value: Value| {
1671            base64::engine::general_purpose::URL_SAFE_NO_PAD
1672                .encode(serde_json::to_vec(&value).unwrap())
1673        };
1674        let access_token = format!(
1675            "{}.{}.signature-is-never-checked",
1676            segment(serde_json::json!({ "alg": "RS256", "typ": "JWT" })),
1677            segment(serde_json::json!({ "exp": expiry })),
1678        );
1679        let body = serde_json::json!({
1680            "auth_mode": "chatgpt",
1681            "tokens": {
1682                "access_token": access_token,
1683                "refresh_token": "refresh",
1684                "id_token": "id",
1685                "account_id": "account",
1686            },
1687            "last_refresh": (now - chrono::TimeDelta::from_std(refreshed_ago).unwrap())
1688                .to_rfc3339(),
1689        });
1690        std::fs::write(home.join("auth.json"), serde_json::to_vec(&body).unwrap()).unwrap();
1691    }
1692
1693    #[cfg(unix)]
1694    fn codex_request_log(log: &Path) -> Vec<Value> {
1695        std::fs::read_to_string(log)
1696            .unwrap()
1697            .lines()
1698            .map(|line| serde_json::from_str::<Value>(line).unwrap())
1699            .collect()
1700    }
1701
1702    #[cfg(unix)]
1703    async fn poll_codex_profile(
1704        directory: &Path,
1705        environment: BTreeMap<String, String>,
1706    ) -> QuotaRefreshOutcome {
1707        let (outcome, client) = refresh_profile(
1708            QuotaRefreshRequest {
1709                profile_id: "codex".into(),
1710                harness: HarnessKind::Codex,
1711                source_home: directory.to_path_buf(),
1712                environment,
1713                cwd: directory.to_path_buf(),
1714            },
1715            None,
1716        )
1717        .await;
1718        if let Some(client) = client {
1719            client.shutdown().await;
1720        }
1721        outcome
1722    }
1723
1724    #[cfg(unix)]
1725    #[tokio::test]
1726    async fn a_codex_login_near_expiry_is_rotated_before_the_usage_query() {
1727        let directory = tempfile::tempdir().unwrap();
1728        // Ten minutes left on a one-hour token: inside the one-hour margin.
1729        write_codex_auth(
1730            directory.path(),
1731            Duration::from_secs(600),
1732            Duration::from_secs(3_000),
1733        );
1734        let (environment, log) = fake_codex_app_server(
1735            directory.path(),
1736            r#"#!/bin/sh
1737read_and_log() {
1738    IFS= read -r line || exit 1
1739    printf '%s\n' "$line" >> "$CODEX_USAGE_TEST_LOG"
1740}
1741read_and_log
1742printf '%s\n' '{"id":1,"result":{}}'
1743read_and_log
1744read_and_log
1745printf '%s\n' '{"auth_mode":"chatgpt","tokens":{"access_token":"rotated"}}' > "$CODEX_AUTH_FILE"
1746printf '%s\n' '{"id":2,"result":{"account":{"type":"chatgpt"}}}'
1747read_and_log
1748printf '%s\n' '{"id":3,"result":{"account":{"type":"chatgpt"}}}'
1749read_and_log
1750printf '%s\n' '{"id":4,"result":{"rateLimits":{"primary":{"usedPercent":25,"windowDurationMins":300}}}}'
1751"#,
1752        );
1753
1754        let outcome = poll_codex_profile(directory.path(), environment).await;
1755
1756        assert_eq!(outcome.report.error, None);
1757        assert_eq!(
1758            outcome.report.five_hour_window().unwrap().remaining_percent,
1759            Some(75)
1760        );
1761        // The rotated file has to reach live sessions, which is what the
1762        // changed-credentials flag asks the daemon to do.
1763        assert!(outcome.credentials_changed);
1764
1765        let messages = codex_request_log(&log);
1766        assert_eq!(messages.len(), 5);
1767        assert_eq!(messages[0]["method"], "initialize");
1768        assert_eq!(messages[1]["method"], "initialized");
1769        assert_eq!(messages[2]["method"], "account/read");
1770        assert_eq!(messages[2]["params"]["refreshToken"], true);
1771        assert_eq!(messages[3]["method"], "account/read");
1772        assert_eq!(messages[3]["params"]["refreshToken"], false);
1773        assert_eq!(messages[4]["method"], "account/rateLimits/read");
1774    }
1775
1776    #[cfg(unix)]
1777    #[tokio::test]
1778    async fn a_codex_login_far_from_expiry_is_polled_without_a_rotation() {
1779        let directory = tempfile::tempdir().unwrap();
1780        // Ten hours left on an eleven-hour token: outside both margins.
1781        write_codex_auth(
1782            directory.path(),
1783            Duration::from_secs(10 * 3_600),
1784            Duration::from_secs(3_600),
1785        );
1786        let (environment, log) = fake_codex_app_server(
1787            directory.path(),
1788            r#"#!/bin/sh
1789read_and_log() {
1790    IFS= read -r line || exit 1
1791    printf '%s\n' "$line" >> "$CODEX_USAGE_TEST_LOG"
1792}
1793read_and_log
1794printf '%s\n' '{"id":1,"result":{}}'
1795read_and_log
1796read_and_log
1797printf '%s\n' '{"id":2,"result":{"account":{"type":"chatgpt"}}}'
1798read_and_log
1799printf '%s\n' '{"id":3,"result":{"rateLimits":{"primary":{"usedPercent":25,"windowDurationMins":300}}}}'
1800"#,
1801        );
1802
1803        let outcome = poll_codex_profile(directory.path(), environment).await;
1804
1805        assert_eq!(outcome.report.error, None);
1806        assert!(!outcome.credentials_changed);
1807
1808        let messages = codex_request_log(&log);
1809        assert_eq!(messages.len(), 4);
1810        assert_eq!(messages[0]["method"], "initialize");
1811        assert_eq!(messages[1]["method"], "initialized");
1812        assert_eq!(messages[2]["params"]["refreshToken"], false);
1813        assert_eq!(messages[3]["method"], "account/rateLimits/read");
1814    }
1815
1816    #[cfg(unix)]
1817    #[tokio::test]
1818    async fn a_codex_app_server_without_the_refresh_flag_still_reports_quota() {
1819        let directory = tempfile::tempdir().unwrap();
1820        write_codex_auth(
1821            directory.path(),
1822            Duration::from_secs(600),
1823            Duration::from_secs(3_000),
1824        );
1825        let (environment, _log) = fake_codex_app_server(
1826            directory.path(),
1827            r#"#!/bin/sh
1828IFS= read -r line || exit 1
1829printf '%s\n' '{"id":1,"result":{}}'
1830IFS= read -r line || exit 1
1831IFS= read -r line || exit 1
1832printf '%s\n' '{"id":2,"error":{"code":-32601,"message":"unknown parameter"}}'
1833IFS= read -r line || exit 1
1834printf '%s\n' '{"id":3,"result":{"account":{"type":"chatgpt"}}}'
1835IFS= read -r line || exit 1
1836printf '%s\n' '{"id":4,"result":{"rateLimits":{"primary":{"usedPercent":40,"windowDurationMins":300}}}}'
1837"#,
1838        );
1839
1840        let outcome = poll_codex_profile(directory.path(), environment).await;
1841
1842        assert_eq!(outcome.report.error, None);
1843        assert_eq!(
1844            outcome.report.five_hour_window().unwrap().remaining_percent,
1845            Some(60)
1846        );
1847    }
1848
1849    #[test]
1850    fn a_codex_refresh_margin_is_an_hour_or_a_tenth_of_the_token_life() {
1851        let hour = 3_600_000;
1852        let now = 1_800_000_000_000;
1853        // A short-lived token: the flat hour decides.
1854        assert!(codex_login_needs_refresh(
1855            Some(now + hour / 2),
1856            Some(now - hour / 2),
1857            now
1858        ));
1859        assert!(!codex_login_needs_refresh(
1860            Some(now + 2 * hour),
1861            Some(now - hour),
1862            now
1863        ));
1864        // A long-lived token: a tenth of its life is wider than the hour.
1865        assert!(codex_login_needs_refresh(
1866            Some(now + 3 * hour),
1867            Some(now - 40 * hour),
1868            now
1869        ));
1870        // Without a last refresh, only the flat hour is known.
1871        assert!(codex_login_needs_refresh(Some(now + hour / 2), None, now));
1872        assert!(!codex_login_needs_refresh(Some(now + 3 * hour), None, now));
1873        // An unreadable expiry is not a reason to spend the refresh token.
1874        assert!(!codex_login_needs_refresh(None, Some(now - hour), now));
1875    }
1876
1877    #[tokio::test]
1878    async fn a_missing_codex_credential_file_asks_for_no_rotation() {
1879        let directory = tempfile::tempdir().unwrap();
1880        assert!(!codex_login_is_near_expiry(&directory.path().join("auth.json")).await);
1881    }
1882
1883    #[tokio::test]
1884    async fn an_unreachable_grok_reports_the_failure_instead_of_a_zero_reading() {
1885        let directory = tempfile::tempdir().unwrap();
1886
1887        let (outcome, _) = refresh_profile(
1888            QuotaRefreshRequest {
1889                profile_id: "grok".into(),
1890                harness: HarnessKind::Grok,
1891                source_home: directory.path().to_path_buf(),
1892                environment: BTreeMap::from([(
1893                    "PATH".to_owned(),
1894                    directory.path().to_string_lossy().into_owned(),
1895                )]),
1896                cwd: directory.path().to_path_buf(),
1897            },
1898            None,
1899        )
1900        .await;
1901        let report = outcome.report;
1902
1903        assert!(report.windows.is_empty());
1904        assert_eq!(
1905            report.error.as_deref(),
1906            Some("Grok Build executable not found")
1907        );
1908    }
1909
1910    #[tokio::test]
1911    async fn deepseek_reports_api_instead_of_inventing_quota() {
1912        let directory = tempfile::tempdir().unwrap();
1913        let (outcome, _) = refresh_profile(
1914            QuotaRefreshRequest {
1915                profile_id: "deepseek".into(),
1916                harness: HarnessKind::Deepseek,
1917                source_home: directory.path().to_path_buf(),
1918                environment: BTreeMap::new(),
1919                cwd: directory.path().to_path_buf(),
1920            },
1921            None,
1922        )
1923        .await;
1924
1925        assert!(outcome.report.windows.is_empty());
1926        assert_eq!(outcome.report.error, None);
1927        assert_eq!(outcome.report.extra.as_deref(), Some(API_LABEL));
1928        assert_eq!(outcome.report.compact(), API_LABEL);
1929    }
1930
1931    #[tokio::test]
1932    async fn expired_claude_credentials_report_login_expired() {
1933        let directory = tempfile::tempdir().unwrap();
1934        std::fs::write(
1935            directory.path().join(".credentials.json"),
1936            serde_json::to_vec(&serde_json::json!({
1937                "claudeAiOauth": {
1938                    "accessToken": "sk-ant-oat01-expired",
1939                    "expiresAt": 1,
1940                }
1941            }))
1942            .unwrap(),
1943        )
1944        .unwrap();
1945
1946        let (outcome, _) = refresh_profile(
1947            QuotaRefreshRequest {
1948                profile_id: "claude2".into(),
1949                harness: HarnessKind::Claude,
1950                source_home: directory.path().to_path_buf(),
1951                environment: BTreeMap::new(),
1952                cwd: directory.path().to_path_buf(),
1953            },
1954            None,
1955        )
1956        .await;
1957        let report = outcome.report;
1958
1959        assert!(report.windows.is_empty());
1960        assert_eq!(report.error.as_deref(), Some(claude_usage::LOGIN_EXPIRED));
1961        assert_eq!(report.compact(), claude_usage::LOGIN_EXPIRED);
1962    }
1963
1964    #[test]
1965    fn a_monthly_window_shares_the_long_window_column_with_a_weekly_one() {
1966        for label in ["Week", "Month"] {
1967            let report = ProfileQuota {
1968                profile_id: "grok".into(),
1969                harness: HarnessKind::Grok,
1970                windows: vec![QuotaWindow {
1971                    label: label.into(),
1972                    remaining_percent: Some(60),
1973                    used: None,
1974                    limit: None,
1975                    resets: None,
1976                    resets_at_epoch_seconds: None,
1977                }],
1978                extra: None,
1979                error: None,
1980                refreshed_at_epoch_seconds: 0,
1981            };
1982
1983            assert!(report.weekly_window().is_some(), "{label}");
1984            assert_eq!(report.compact(), format!("{label} 60% left"));
1985        }
1986    }
1987
1988    #[test]
1989    fn kimi_uses_percent_left_and_hides_a_short_window_on_sustainable_pace() {
1990        let report = ProfileQuota {
1991            profile_id: "kimi".into(),
1992            harness: HarnessKind::Kimi,
1993            windows: vec![
1994                QuotaWindow {
1995                    label: "Week".into(),
1996                    remaining_percent: Some(94),
1997                    used: Some(6),
1998                    limit: Some(100),
1999                    resets: Some("12:22 Aug 18".into()),
2000                    resets_at_epoch_seconds: Some(604_800),
2001                },
2002                QuotaWindow {
2003                    label: "5H".into(),
2004                    remaining_percent: Some(97),
2005                    used: Some(3),
2006                    limit: Some(100),
2007                    resets: Some("10:22 Aug 13".into()),
2008                    resets_at_epoch_seconds: Some(18_000),
2009                },
2010            ],
2011            extra: None,
2012            error: None,
2013            refreshed_at_epoch_seconds: 3_600,
2014        };
2015
2016        assert_eq!(report.compact(), "Week 94% left, resets 12:22 Aug 18");
2017    }
2018
2019    #[test]
2020    fn short_window_is_shown_only_when_burn_rate_projects_early_exhaustion() {
2021        let window = QuotaWindow {
2022            label: "5H".into(),
2023            remaining_percent: Some(70),
2024            used: None,
2025            limit: None,
2026            resets: Some("later".into()),
2027            resets_at_epoch_seconds: Some(14_400),
2028        };
2029        assert!(projects_exhaustion_before_reset(&window, 0));
2030
2031        let sustainable = QuotaWindow {
2032            remaining_percent: Some(80),
2033            ..window
2034        };
2035        assert!(!projects_exhaustion_before_reset(&sustainable, 0));
2036    }
2037
2038    #[derive(Clone, Default)]
2039    struct KimiServerState {
2040        refresh_forms: Arc<Mutex<Vec<String>>>,
2041    }
2042
2043    async fn test_kimi_usage(headers: HeaderMap) -> (StatusCode, Json<Value>) {
2044        let accepted = headers
2045            .get(reqwest::header::AUTHORIZATION)
2046            .and_then(|value| value.to_str().ok())
2047            == Some("Bearer fresh-access");
2048        if accepted {
2049            (
2050                StatusCode::OK,
2051                Json(serde_json::json!({
2052                    "usage": {"name": "Weekly", "used": 1, "limit": 100}
2053                })),
2054            )
2055        } else {
2056            (StatusCode::UNAUTHORIZED, Json(serde_json::json!({})))
2057        }
2058    }
2059
2060    async fn test_kimi_refresh(State(state): State<KimiServerState>, body: Bytes) -> Json<Value> {
2061        state
2062            .refresh_forms
2063            .lock()
2064            .unwrap()
2065            .push(String::from_utf8(body.to_vec()).unwrap());
2066        Json(serde_json::json!({
2067            "access_token": "fresh-access",
2068            "refresh_token": "fresh-refresh",
2069            "expires_in": 900,
2070            "scope": "kimi-code",
2071            "token_type": "Bearer"
2072        }))
2073    }
2074
2075    #[tokio::test]
2076    async fn kimi_quota_refreshes_after_unauthorized_and_retries() {
2077        let state = KimiServerState::default();
2078        let app = Router::new()
2079            .route("/coding/v1/usages", get(test_kimi_usage))
2080            .route("/api/oauth/token", post(test_kimi_refresh))
2081            .with_state(state.clone());
2082        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2083        let address = listener.local_addr().unwrap();
2084        let server = tokio::spawn(async move {
2085            axum::serve(listener, app).await.unwrap();
2086        });
2087
2088        let home = tempfile::tempdir().unwrap();
2089        let credentials_path = home.path().join("credentials/kimi-code.json");
2090        tokio::fs::create_dir_all(credentials_path.parent().unwrap())
2091            .await
2092            .unwrap();
2093        let future_expiry = SystemTime::now()
2094            .duration_since(SystemTime::UNIX_EPOCH)
2095            .unwrap()
2096            .as_secs()
2097            + 3_600;
2098        tokio::fs::write(
2099            &credentials_path,
2100            serde_json::to_vec(&serde_json::json!({
2101                "access_token": "rejected-access",
2102                "refresh_token": "old-refresh",
2103                "expires_at": future_expiry,
2104                "scope": "kimi-code",
2105                "token_type": "Bearer",
2106                "expires_in": 900
2107            }))
2108            .unwrap(),
2109        )
2110        .await
2111        .unwrap();
2112        let endpoint = format!("http://{address}");
2113        let environment = HashMap::from([
2114            ("KIMI_CODE_BASE_URL".into(), format!("{endpoint}/coding/v1")),
2115            ("KIMI_CODE_OAUTH_HOST".into(), endpoint),
2116        ]);
2117
2118        let (windows, _) = query_kimi(home.path(), &environment).await.unwrap();
2119
2120        assert_eq!(windows[0].used, Some(1));
2121        let form = {
2122            let forms = state.refresh_forms.lock().unwrap();
2123            assert_eq!(forms.len(), 1);
2124            url::form_urlencoded::parse(forms[0].as_bytes())
2125                .into_owned()
2126                .collect::<HashMap<_, _>>()
2127        };
2128        assert_eq!(
2129            form.get("grant_type").map(String::as_str),
2130            Some("refresh_token")
2131        );
2132        assert_eq!(
2133            form.get("refresh_token").map(String::as_str),
2134            Some("old-refresh")
2135        );
2136        let saved = read_kimi_credentials(&credentials_path).await.unwrap();
2137        assert_eq!(saved.access_token, "fresh-access");
2138        assert_eq!(saved.refresh_token, "fresh-refresh");
2139        assert!(!home.path().join("oauth/kimi-code.lock").exists());
2140        server.abort();
2141    }
2142
2143    /// Backdate the lock directory the way a holder that stopped heartbeating
2144    /// leaves it behind.
2145    fn age_kimi_lock(path: &Path, age: Duration) {
2146        touch_kimi_lock(path, SystemTime::now() - age).expect("backdate lock directory");
2147    }
2148
2149    #[tokio::test]
2150    async fn a_kimi_refresh_lock_left_by_a_crashed_holder_is_broken_and_reacquired() {
2151        let home = tempfile::tempdir().unwrap();
2152        let lock = home.path().join("oauth/kimi-code.lock");
2153        std::fs::create_dir_all(&lock).unwrap();
2154        age_kimi_lock(&lock, KIMI_LOCK_STALE_AFTER + Duration::from_secs(60));
2155
2156        let started = std::time::Instant::now();
2157        let held = KimiRefreshLock::acquire_within(home.path(), Duration::from_secs(10))
2158            .await
2159            .expect("an orphaned lock must not block a refresh");
2160        let waited = started.elapsed();
2161
2162        assert!(
2163            waited < Duration::from_secs(5),
2164            "acquisition waited {waited:?}"
2165        );
2166        held.release().await.expect("release an uncontested lock");
2167        assert!(!lock.exists(), "the released lock must be gone");
2168    }
2169
2170    #[tokio::test]
2171    async fn a_heartbeating_kimi_refresh_lock_is_not_broken_by_a_waiter() {
2172        let home = tempfile::tempdir().unwrap();
2173        let lock = home.path().join("oauth/kimi-code.lock");
2174        std::fs::create_dir_all(&lock).unwrap();
2175
2176        let error = KimiRefreshLock::acquire_within(home.path(), Duration::from_millis(600))
2177            .await
2178            .err()
2179            .expect("a lock with a live holder must be waited out, not stolen");
2180
2181        assert!(
2182            error.to_string().contains("kimi-code.lock"),
2183            "the timeout must name the lock: {error}"
2184        );
2185        assert!(lock.exists(), "a live holder's lock must survive a waiter");
2186    }
2187
2188    /// The Kimi Code CLI breaks a lock whose modification time is more than
2189    /// `KIMI_CLI_LOCK_STALE_AFTER` old, so Mjolnir's beats have to be frequent
2190    /// enough that a stalled heartbeat task still cannot cost it a live lock.
2191    #[tokio::test]
2192    async fn a_held_kimi_lock_republishes_its_mtime_several_times_per_cli_break_window() {
2193        let home = tempfile::tempdir().unwrap();
2194        let held = KimiRefreshLock::acquire(home.path()).await.unwrap();
2195        let lock = home.path().join("oauth/kimi-code.lock");
2196
2197        // Half the peer's break window: two beats have to land inside it, so
2198        // Mjolnir publishes at least four times per window and can miss several in
2199        // a row and still hold the lock.
2200        let watched = KIMI_CLI_LOCK_STALE_AFTER / 2;
2201        let deadline = tokio::time::Instant::now() + watched;
2202        let mut published = vec![kimi_lock_mtime(&lock).unwrap().expect("the created lock")];
2203        while tokio::time::Instant::now() < deadline {
2204            tokio::time::sleep(Duration::from_millis(50)).await;
2205            let observed = kimi_lock_mtime(&lock).unwrap().expect("a held lock");
2206            if published.last() != Some(&observed) {
2207                published.push(observed);
2208            }
2209            assert_eq!(
2210                std::fs::read_dir(&lock).unwrap().count(),
2211                0,
2212                "the CLI releases this lock with a plain rmdir, so it must stay empty"
2213            );
2214        }
2215
2216        assert!(
2217            published.len() >= 3,
2218            "the lock's modification time moved {} times in {watched:?}; the Kimi Code CLI breaks a lock after {KIMI_CLI_LOCK_STALE_AFTER:?} without a beat",
2219            published.len() - 1
2220        );
2221        held.release().await.expect("release an uncontested lock");
2222    }
2223
2224    #[tokio::test]
2225    async fn a_stolen_kimi_refresh_lock_is_left_to_its_new_holder_and_fails_the_refresh() {
2226        let home = tempfile::tempdir().unwrap();
2227        let held = KimiRefreshLock::acquire(home.path()).await.unwrap();
2228        let lock = home.path().join("oauth/kimi-code.lock");
2229
2230        // The Kimi Code CLI breaks a lock it judges stale and takes it over:
2231        // rmdir, mkdir, then its own modification time. The CLI dates those
2232        // ahead of the clock; date this one far enough ahead that only the new
2233        // holder could have written it.
2234        std::fs::remove_dir(&lock).unwrap();
2235        std::fs::create_dir(&lock).unwrap();
2236        let thief = SystemTime::now() + Duration::from_secs(30);
2237        touch_kimi_lock(&lock, thief).unwrap();
2238
2239        // Mjolnir must stop beating: a touch on the CLI's lock trips the CLI's own
2240        // mtime ownership check and it abandons its refresh with ECOMPROMISED.
2241        tokio::time::sleep(2 * KIMI_LOCK_HEARTBEAT_INTERVAL + Duration::from_millis(400)).await;
2242        let observed = kimi_lock_mtime(&lock)
2243            .unwrap()
2244            .expect("the new holder's lock");
2245        assert!(
2246            observed.duration_since(SystemTime::now()).is_ok(),
2247            "Mjolnir kept heartbeating a lock it no longer holds: the directory carries {} instead of the new holder's {}",
2248            epoch_label(observed),
2249            epoch_label(thief)
2250        );
2251
2252        let error = held
2253            .release()
2254            .await
2255            .expect_err("a refresh that lost its lock must fail loudly");
2256        let message = error.to_string();
2257        assert!(message.contains("kimi-code.lock"), "{message}");
2258        assert!(message.contains("another process took it"), "{message}");
2259        assert!(
2260            lock.exists(),
2261            "Mjolnir must not remove a lock another holder owns"
2262        );
2263    }
2264
2265    #[tokio::test]
2266    async fn a_kimi_refresh_lock_removed_underneath_hel_fails_the_refresh() {
2267        let home = tempfile::tempdir().unwrap();
2268        let held = KimiRefreshLock::acquire(home.path()).await.unwrap();
2269        let lock = home.path().join("oauth/kimi-code.lock");
2270        std::fs::remove_dir(&lock).unwrap();
2271
2272        let error = held
2273            .release()
2274            .await
2275            .expect_err("a refresh that lost its lock must fail loudly");
2276
2277        assert!(
2278            error.to_string().contains("another process removed it"),
2279            "{error}"
2280        );
2281        assert!(!lock.exists(), "Mjolnir must not recreate a lock it lost");
2282    }
2283
2284    fn kimi_pair(access: &str, refresh: &str, expires_at: i64) -> KimiCredentials {
2285        KimiCredentials {
2286            access_token: access.into(),
2287            refresh_token: refresh.into(),
2288            expires_at,
2289            scope: "kimi-code".into(),
2290            token_type: "Bearer".into(),
2291            expires_in: 900,
2292        }
2293    }
2294
2295    fn a_stolen_lock() -> KimiLockLoss {
2296        KimiLockLoss::Stolen {
2297            published: SystemTime::UNIX_EPOCH,
2298            observed: SystemTime::UNIX_EPOCH + Duration::from_secs(30),
2299        }
2300    }
2301
2302    #[test]
2303    fn a_lock_still_held_saves_the_refreshed_pair_and_releases() {
2304        let active = kimi_pair("spent-access", "spent-refresh", 100);
2305
2306        assert_eq!(
2307            decide_kimi_refresh_persist(&Ok(SystemTime::now()), None, &active),
2308            KimiRefreshPersist::Save
2309        );
2310    }
2311
2312    #[test]
2313    fn an_unprovable_ownership_check_still_saves_the_refreshed_pair() {
2314        let active = kimi_pair("spent-access", "spent-refresh", 100);
2315        let unproven = Err(KimiLockLoss::Unproven("stat failed".into()));
2316
2317        assert_eq!(
2318            decide_kimi_refresh_persist(&unproven, None, &active),
2319            KimiRefreshPersist::Save,
2320            "a failed stat proves nothing about the lock and must not discard valid tokens"
2321        );
2322    }
2323
2324    #[test]
2325    fn a_stolen_lock_adopts_the_thiefs_newer_credentials() {
2326        let active = kimi_pair("spent-access", "spent-refresh", 100);
2327        let thiefs = kimi_pair("thief-access", "thief-refresh", 900);
2328
2329        assert_eq!(
2330            decide_kimi_refresh_persist(&Err(a_stolen_lock()), Some(&thiefs), &active),
2331            KimiRefreshPersist::Adopt {
2332                access_token: "thief-access".into(),
2333                loss: a_stolen_lock(),
2334            }
2335        );
2336    }
2337
2338    #[test]
2339    fn a_removed_lock_adopts_the_newer_credentials_left_on_disk() {
2340        let active = kimi_pair("spent-access", "spent-refresh", 100);
2341        let thiefs = kimi_pair("thief-access", "thief-refresh", 900);
2342
2343        assert_eq!(
2344            decide_kimi_refresh_persist(&Err(KimiLockLoss::Gone), Some(&thiefs), &active),
2345            KimiRefreshPersist::Adopt {
2346                access_token: "thief-access".into(),
2347                loss: KimiLockLoss::Gone,
2348            }
2349        );
2350    }
2351
2352    /// The pair on disk is dead whoever wrote it: its refresh token is the one
2353    /// this refresh just spent at the server.
2354    #[test]
2355    fn a_stolen_lock_saves_hels_pair_over_the_spent_one_on_disk() {
2356        let active = kimi_pair("spent-access", "spent-refresh", 100);
2357        let on_disk = active.clone();
2358
2359        assert_eq!(
2360            decide_kimi_refresh_persist(&Err(a_stolen_lock()), Some(&on_disk), &active),
2361            KimiRefreshPersist::SaveContested(a_stolen_lock())
2362        );
2363    }
2364
2365    #[test]
2366    fn an_unreadable_credentials_file_after_a_lost_lock_saves_hels_pair() {
2367        let active = kimi_pair("spent-access", "spent-refresh", 100);
2368
2369        assert_eq!(
2370            decide_kimi_refresh_persist(&Err(KimiLockLoss::Gone), None, &active),
2371            KimiRefreshPersist::SaveContested(KimiLockLoss::Gone),
2372            "nothing readable is newer, so the refreshed pair is the only live one"
2373        );
2374    }
2375
2376    #[test]
2377    fn any_rotated_field_marks_the_disk_pair_as_the_other_refreshers() {
2378        let active = kimi_pair("spent-access", "spent-refresh", 100);
2379        let rotations = [
2380            kimi_pair("other-access", "spent-refresh", 100),
2381            kimi_pair("spent-access", "other-refresh", 100),
2382            kimi_pair("spent-access", "spent-refresh", 900),
2383        ];
2384
2385        for on_disk in &rotations {
2386            assert!(
2387                matches!(
2388                    decide_kimi_refresh_persist(&Err(a_stolen_lock()), Some(on_disk), &active),
2389                    KimiRefreshPersist::Adopt { .. }
2390                ),
2391                "{on_disk:?} is a rotated pair, not the spent one"
2392            );
2393        }
2394    }
2395
2396    #[test]
2397    fn a_disk_pair_that_differs_only_in_description_is_still_the_spent_one() {
2398        let active = kimi_pair("spent-access", "spent-refresh", 100);
2399        let on_disk = KimiCredentials {
2400            scope: "kimi-code extra".into(),
2401            token_type: "bearer".into(),
2402            expires_in: 1_800,
2403            ..active.clone()
2404        };
2405
2406        assert_eq!(
2407            decide_kimi_refresh_persist(&Err(a_stolen_lock()), Some(&on_disk), &active),
2408            KimiRefreshPersist::SaveContested(a_stolen_lock())
2409        );
2410    }
2411
2412    #[derive(Clone)]
2413    struct KimiThiefState {
2414        home: std::path::PathBuf,
2415        /// The pair the lock's new holder stores before Mjolnir's own refresh
2416        /// returns, when it got that far.
2417        winner: Option<Value>,
2418    }
2419
2420    /// Answer the refresh, but take the lock over first the way the Kimi Code
2421    /// CLI takes over one it judges stale: rmdir, mkdir, then a modification
2422    /// time of its own, dated ahead of the clock as the CLI dates its locks.
2423    async fn test_kimi_refresh_stealing_the_lock(
2424        State(state): State<KimiThiefState>,
2425        _body: Bytes,
2426    ) -> Json<Value> {
2427        let lock = state.home.join("oauth/kimi-code.lock");
2428        std::fs::remove_dir(&lock).unwrap();
2429        std::fs::create_dir(&lock).unwrap();
2430        touch_kimi_lock(&lock, SystemTime::now() + Duration::from_secs(30)).unwrap();
2431        if let Some(winner) = &state.winner {
2432            std::fs::write(
2433                state.home.join("credentials/kimi-code.json"),
2434                serde_json::to_vec(winner).unwrap(),
2435            )
2436            .unwrap();
2437        }
2438        Json(serde_json::json!({
2439            "access_token": "hel-access",
2440            "refresh_token": "hel-refresh",
2441            "expires_in": 900,
2442            "scope": "kimi-code",
2443            "token_type": "Bearer"
2444        }))
2445    }
2446
2447    /// Serve one refresh that steals the lock while it is in flight, and hand
2448    /// back what `ensure_fresh_kimi_token` made of it.
2449    async fn refresh_against_a_lock_thief(
2450        home: &Path,
2451        winner: Option<Value>,
2452    ) -> (Result<String>, std::path::PathBuf) {
2453        let credentials_path = home.join("credentials/kimi-code.json");
2454        tokio::fs::create_dir_all(credentials_path.parent().unwrap())
2455            .await
2456            .unwrap();
2457        let soon = SystemTime::now()
2458            .duration_since(SystemTime::UNIX_EPOCH)
2459            .unwrap()
2460            .as_secs() as i64
2461            + 10;
2462        tokio::fs::write(
2463            &credentials_path,
2464            serde_json::to_vec(&serde_json::json!({
2465                "access_token": "spent-access",
2466                "refresh_token": "spent-refresh",
2467                "expires_at": soon,
2468                "scope": "kimi-code",
2469                "token_type": "Bearer",
2470                "expires_in": 900
2471            }))
2472            .unwrap(),
2473        )
2474        .await
2475        .unwrap();
2476
2477        let app = Router::new()
2478            .route(
2479                "/api/oauth/token",
2480                post(test_kimi_refresh_stealing_the_lock),
2481            )
2482            .with_state(KimiThiefState {
2483                home: home.to_path_buf(),
2484                winner,
2485            });
2486        let listener = tokio::net::TcpListener::bind("127.0.0.1:0").await.unwrap();
2487        let address = listener.local_addr().unwrap();
2488        let server = tokio::spawn(async move {
2489            axum::serve(listener, app).await.unwrap();
2490        });
2491        let environment =
2492            HashMap::from([("KIMI_CODE_OAUTH_HOST".into(), format!("http://{address}"))]);
2493
2494        let token = ensure_fresh_kimi_token(
2495            &reqwest::Client::new(),
2496            home,
2497            &credentials_path,
2498            &environment,
2499            false,
2500            None,
2501        )
2502        .await;
2503        server.abort();
2504        (token, credentials_path)
2505    }
2506
2507    #[tokio::test]
2508    async fn a_refresh_that_loses_its_lock_returns_the_winners_stored_token() {
2509        let home = tempfile::tempdir().unwrap();
2510        let winner = serde_json::json!({
2511            "access_token": "winner-access",
2512            "refresh_token": "winner-refresh",
2513            "expires_at": 4_102_444_800i64,
2514            "scope": "kimi-code",
2515            "token_type": "Bearer",
2516            "expires_in": 900
2517        });
2518
2519        let (token, credentials_path) =
2520            refresh_against_a_lock_thief(home.path(), Some(winner)).await;
2521
2522        assert_eq!(
2523            token.unwrap(),
2524            "winner-access",
2525            "a contested lock must not fail a refresh when a live token exists"
2526        );
2527        let saved = read_kimi_credentials(&credentials_path).await.unwrap();
2528        assert_eq!(
2529            saved.access_token, "winner-access",
2530            "Mjolnir must not clobber the credentials the lock's new holder stored"
2531        );
2532        assert!(
2533            home.path().join("oauth/kimi-code.lock").exists(),
2534            "Mjolnir must not remove a lock another holder owns"
2535        );
2536    }
2537
2538    /// The pair the thief left behind is the one this refresh already spent, so
2539    /// it is dead: only Mjolnir's pair can still authenticate, and storing it is
2540    /// what keeps the peer's own recovery re-read working.
2541    #[tokio::test]
2542    async fn a_refresh_that_loses_its_lock_saves_its_pair_over_the_spent_one() {
2543        let home = tempfile::tempdir().unwrap();
2544
2545        let (token, credentials_path) = refresh_against_a_lock_thief(home.path(), None).await;
2546
2547        assert_eq!(token.unwrap(), "hel-access");
2548        let saved = read_kimi_credentials(&credentials_path).await.unwrap();
2549        assert_eq!(
2550            saved.access_token, "hel-access",
2551            "leaving the spent pair on disk would force a `kimi login`"
2552        );
2553        assert_eq!(saved.refresh_token, "hel-refresh");
2554        assert!(
2555            home.path().join("oauth/kimi-code.lock").exists(),
2556            "Mjolnir must not remove a lock another holder owns"
2557        );
2558    }
2559
2560    /// The child is reaped by the shutdown, so a live pid means it was never
2561    /// stopped.
2562    #[cfg(unix)]
2563    fn process_is_gone(pid: i32) -> bool {
2564        // SAFETY: signal 0 only probes whether the process exists.
2565        unsafe { libc::kill(pid, 0) != 0 }
2566    }
2567
2568    #[cfg(unix)]
2569    #[tokio::test]
2570    async fn dropping_a_profile_from_the_configuration_stops_its_codex_quota_client() {
2571        use std::os::unix::fs::PermissionsExt;
2572
2573        let directory = tempfile::tempdir().unwrap();
2574        let executable = directory.path().join("codex");
2575        let pid_file = directory.path().join("codex.pid");
2576        // A `codex app-server` stand-in: answer one quota refresh, then stay
2577        // alive on stdin the way the real one does between refreshes.
2578        std::fs::write(
2579            &executable,
2580            r#"#!/bin/sh
2581printf '%s\n' "$$" > "$CODEX_QUOTA_TEST_PID"
2582IFS= read -r line || exit 0
2583printf '%s\n' '{"id":1,"result":{}}'
2584IFS= read -r line || exit 0
2585IFS= read -r line || exit 0
2586printf '%s\n' '{"id":2,"result":{"account":{"type":"chatgpt"}}}'
2587IFS= read -r line || exit 0
2588printf '%s\n' '{"id":3,"result":{"rateLimits":{"primary":{"usedPercent":25,"windowDurationMins":300}}}}'
2589while IFS= read -r line; do :; done
2590"#,
2591        )
2592        .unwrap();
2593        std::fs::set_permissions(&executable, std::fs::Permissions::from_mode(0o755)).unwrap();
2594        let request = QuotaRefreshRequest {
2595            profile_id: "codex-1".into(),
2596            harness: HarnessKind::Codex,
2597            source_home: directory.path().to_path_buf(),
2598            environment: BTreeMap::from([
2599                (
2600                    "PATH".to_owned(),
2601                    directory.path().to_string_lossy().into_owned(),
2602                ),
2603                (
2604                    "CODEX_QUOTA_TEST_PID".to_owned(),
2605                    pid_file.to_string_lossy().into_owned(),
2606                ),
2607            ]),
2608            cwd: directory.path().to_path_buf(),
2609        };
2610
2611        let mut quotas = QuotaManager::default();
2612        quotas.refresh_profiles(vec![request], |_| async {}).await;
2613
2614        assert_eq!(
2615            quotas.reports()["codex-1"].error,
2616            None,
2617            "the stand-in app-server must answer the quota query"
2618        );
2619        let pid = std::fs::read_to_string(&pid_file)
2620            .unwrap()
2621            .trim()
2622            .parse::<i32>()
2623            .unwrap();
2624        assert!(
2625            !process_is_gone(pid),
2626            "the app-server child is cached between refreshes"
2627        );
2628
2629        // The profile leaves the configuration, so the next batch no longer
2630        // carries it.
2631        quotas.refresh_profiles(Vec::new(), |_| async {}).await;
2632
2633        assert!(
2634            process_is_gone(pid),
2635            "a profile removed from the configuration must not leave its `codex app-server` child running"
2636        );
2637        quotas.shutdown().await;
2638    }
2639
2640    #[test]
2641    fn reset_time_normalization_uses_24_hour_month_day_format() {
2642        let paris = FixedOffset::east_opt(2 * 3_600).expect("offset");
2643        let reset = paris
2644            .with_ymd_and_hms(2026, 6, 17, 16, 49, 0)
2645            .single()
2646            .expect("instant");
2647        assert_eq!(format_reset_label(reset), "16:49 Jun 17");
2648        assert_eq!(
2649            normalize_reset_text("Jun 17 at 4:49pm").as_deref(),
2650            Some("16:49 Jun 17")
2651        );
2652    }
2653
2654    #[test]
2655    fn reset_timestamp_accepts_seconds_and_milliseconds() {
2656        let seconds = 1_781_712_540_f64;
2657        assert_eq!(
2658            format_reset_local(seconds),
2659            format_reset_local(seconds * 1_000.0)
2660        );
2661        assert_eq!(
2662            format_reset_local(seconds),
2663            format_reset_local_seconds(seconds as i64)
2664        );
2665    }
2666
2667    #[test]
2668    fn time_only_reset_is_rendered_as_the_next_datetime() {
2669        let zone = FixedOffset::west_opt(5 * 3_600).expect("offset");
2670        let now = zone
2671            .with_ymd_and_hms(2026, 8, 10, 14, 0, 0)
2672            .single()
2673            .expect("now");
2674        assert_eq!(
2675            normalize_reset_at("3:30 PM (America/Chicago)", now)
2676                .map(format_reset_label)
2677                .as_deref(),
2678            Some("15:30 Aug 10")
2679        );
2680        assert_eq!(
2681            normalize_reset_at("at 1pm (America/Chicago)", now)
2682                .map(format_reset_label)
2683                .as_deref(),
2684            Some("13:00 Aug 11")
2685        );
2686    }
2687
2688    #[test]
2689    fn claude_comma_separated_reset_is_normalized() {
2690        let zone = FixedOffset::west_opt(5 * 3_600).expect("offset");
2691        let now = zone
2692            .with_ymd_and_hms(2026, 8, 11, 7, 0, 0)
2693            .single()
2694            .expect("now");
2695        assert_eq!(
2696            normalize_reset_at("Aug 14, 4am (America/Chicago)", now)
2697                .map(format_reset_label)
2698                .as_deref(),
2699            Some("04:00 Aug 14")
2700        );
2701    }
2702}