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