Skip to main content

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