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