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