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