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