Skip to main content

mj_controller/
hel_quota.rs

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