1use anyhow::{Context, Result};
2use std::fs;
3use std::path::{Path, PathBuf};
4
5use crate::auth;
6use crate::fs as codex_fs;
7use crate::paths;
8
9pub struct CacheEntry {
10 pub non_weekly_label: String,
11 pub non_weekly_remaining: i64,
12 pub non_weekly_reset_epoch: Option<i64>,
13 pub weekly_remaining: i64,
14 pub weekly_reset_epoch: i64,
15}
16
17pub fn clear_starship_cache() -> Result<()> {
18 let root = cache_root().context("cache root")?;
19 if !root.is_absolute() {
20 anyhow::bail!(
21 "codex-rate-limits: refusing to clear cache with non-absolute cache root: {}",
22 root.display()
23 );
24 }
25 if root == Path::new("/") {
26 anyhow::bail!(
27 "codex-rate-limits: refusing to clear cache with invalid cache root: {}",
28 root.display()
29 );
30 }
31
32 let cache_dir = root.join("codex").join("starship-rate-limits");
33 let cache_dir_str = cache_dir.to_string_lossy();
34 if !cache_dir_str.ends_with("/codex/starship-rate-limits") {
35 anyhow::bail!(
36 "codex-rate-limits: refusing to clear unexpected cache dir: {}",
37 cache_dir.display()
38 );
39 }
40
41 if cache_dir.is_dir() {
42 fs::remove_dir_all(&cache_dir).ok();
43 }
44
45 Ok(())
46}
47
48pub fn cache_file_for_target(target_file: &Path) -> Result<PathBuf> {
49 let cache_dir = starship_cache_dir().context("cache dir")?;
50
51 if let Some(secret_dir) = paths::resolve_secret_dir() {
52 if target_file.starts_with(&secret_dir) {
53 let display = secret_file_basename(target_file)?;
54 let key = cache_key(&display)?;
55 return Ok(cache_dir.join(format!("{key}.kv")));
56 }
57
58 if let Some(secret_name) = secret_name_for_auth(target_file, &secret_dir) {
59 let key = cache_key(&secret_name)?;
60 return Ok(cache_dir.join(format!("{key}.kv")));
61 }
62 }
63
64 let hash = codex_fs::sha256_file(target_file)?;
65 Ok(cache_dir.join(format!("auth_{}.kv", hash.to_lowercase())))
66}
67
68pub fn secret_name_for_target(target_file: &Path) -> Option<String> {
69 let secret_dir = paths::resolve_secret_dir()?;
70 if target_file.starts_with(&secret_dir) {
71 return secret_file_basename(target_file).ok();
72 }
73 secret_name_for_auth(target_file, &secret_dir)
74}
75
76pub fn read_cache_entry(target_file: &Path) -> Result<CacheEntry> {
77 let cache_file = cache_file_for_target(target_file)?;
78 if !cache_file.is_file() {
79 anyhow::bail!(
80 "codex-rate-limits: cache not found (run codex-rate-limits without --cached, or codex-starship, to populate): {}",
81 cache_file.display()
82 );
83 }
84
85 let content = fs::read_to_string(&cache_file)
86 .with_context(|| format!("failed to read cache: {}", cache_file.display()))?;
87 let mut non_weekly_label: Option<String> = None;
88 let mut non_weekly_remaining: Option<i64> = None;
89 let mut non_weekly_reset_epoch: Option<i64> = None;
90 let mut weekly_remaining: Option<i64> = None;
91 let mut weekly_reset_epoch: Option<i64> = None;
92
93 for line in content.lines() {
94 if let Some(value) = line.strip_prefix("non_weekly_label=") {
95 non_weekly_label = Some(value.to_string());
96 } else if let Some(value) = line.strip_prefix("non_weekly_remaining=") {
97 non_weekly_remaining = value.parse::<i64>().ok();
98 } else if let Some(value) = line.strip_prefix("non_weekly_reset_epoch=") {
99 non_weekly_reset_epoch = value.parse::<i64>().ok();
100 } else if let Some(value) = line.strip_prefix("weekly_remaining=") {
101 weekly_remaining = value.parse::<i64>().ok();
102 } else if let Some(value) = line.strip_prefix("weekly_reset_epoch=") {
103 weekly_reset_epoch = value.parse::<i64>().ok();
104 }
105 }
106
107 let non_weekly_label = match non_weekly_label {
108 Some(value) if !value.is_empty() => value,
109 _ => anyhow::bail!(
110 "codex-rate-limits: invalid cache (missing non-weekly data): {}",
111 cache_file.display()
112 ),
113 };
114 let non_weekly_remaining = match non_weekly_remaining {
115 Some(value) => value,
116 _ => anyhow::bail!(
117 "codex-rate-limits: invalid cache (missing non-weekly data): {}",
118 cache_file.display()
119 ),
120 };
121 let weekly_remaining = match weekly_remaining {
122 Some(value) => value,
123 _ => anyhow::bail!(
124 "codex-rate-limits: invalid cache (missing weekly data): {}",
125 cache_file.display()
126 ),
127 };
128 let weekly_reset_epoch = match weekly_reset_epoch {
129 Some(value) => value,
130 _ => anyhow::bail!(
131 "codex-rate-limits: invalid cache (missing weekly data): {}",
132 cache_file.display()
133 ),
134 };
135
136 Ok(CacheEntry {
137 non_weekly_label,
138 non_weekly_remaining,
139 non_weekly_reset_epoch,
140 weekly_remaining,
141 weekly_reset_epoch,
142 })
143}
144
145pub fn write_starship_cache(
146 target_file: &Path,
147 fetched_at_epoch: i64,
148 non_weekly_label: &str,
149 non_weekly_remaining: i64,
150 weekly_remaining: i64,
151 weekly_reset_epoch: i64,
152 non_weekly_reset_epoch: Option<i64>,
153) -> Result<()> {
154 let cache_file = cache_file_for_target(target_file)?;
155 if let Some(parent) = cache_file.parent() {
156 fs::create_dir_all(parent)?;
157 }
158
159 let mut lines = Vec::new();
160 lines.push(format!("fetched_at={fetched_at_epoch}"));
161 lines.push(format!("non_weekly_label={non_weekly_label}"));
162 lines.push(format!("non_weekly_remaining={non_weekly_remaining}"));
163 if let Some(epoch) = non_weekly_reset_epoch {
164 lines.push(format!("non_weekly_reset_epoch={epoch}"));
165 }
166 lines.push(format!("weekly_remaining={weekly_remaining}"));
167 lines.push(format!("weekly_reset_epoch={weekly_reset_epoch}"));
168
169 let data = lines.join("\n");
170 codex_fs::write_atomic(&cache_file, data.as_bytes(), codex_fs::SECRET_FILE_MODE)?;
171 Ok(())
172}
173
174fn starship_cache_dir() -> Result<PathBuf> {
175 let root = cache_root().context("cache root")?;
176 Ok(root.join("codex").join("starship-rate-limits"))
177}
178
179fn cache_root() -> Option<PathBuf> {
180 if let Ok(path) = std::env::var("ZSH_CACHE_DIR")
181 && !path.is_empty()
182 {
183 return Some(PathBuf::from(path));
184 }
185 let zdotdir = paths::resolve_zdotdir()?;
186 Some(zdotdir.join("cache"))
187}
188
189fn secret_name_for_auth(auth_file: &Path, secret_dir: &Path) -> Option<String> {
190 let auth_key = auth::identity_key_from_auth_file(auth_file)
191 .ok()
192 .flatten()?;
193 let entries = std::fs::read_dir(secret_dir).ok()?;
194 for entry in entries.flatten() {
195 let path = entry.path();
196 if path.extension().and_then(|s| s.to_str()) != Some("json") {
197 continue;
198 }
199 let candidate_key = match auth::identity_key_from_auth_file(&path).ok().flatten() {
200 Some(value) => value,
201 None => continue,
202 };
203 if candidate_key == auth_key {
204 return secret_file_basename(&path).ok();
205 }
206 }
207 None
208}
209
210fn secret_file_basename(path: &Path) -> Result<String> {
211 let file = path
212 .file_name()
213 .and_then(|name| name.to_str())
214 .unwrap_or_default();
215 let base = file.trim_end_matches(".json");
216 Ok(base.to_string())
217}
218
219fn cache_key(name: &str) -> Result<String> {
220 if name.is_empty() {
221 anyhow::bail!("missing cache key name");
222 }
223 let mut key = String::new();
224 for ch in name.to_lowercase().chars() {
225 if ch.is_ascii_alphanumeric() {
226 key.push(ch);
227 } else {
228 key.push('_');
229 }
230 }
231 while key.starts_with('_') {
232 key.remove(0);
233 }
234 while key.ends_with('_') {
235 key.pop();
236 }
237 if key.is_empty() {
238 anyhow::bail!("invalid cache key name");
239 }
240 Ok(key)
241}
242
243#[cfg(test)]
244mod tests {
245 use super::{
246 cache_file_for_target, clear_starship_cache, read_cache_entry, secret_name_for_target,
247 write_starship_cache,
248 };
249 use crate::fs as codex_fs;
250 use nils_test_support::{EnvGuard, GlobalStateLock};
251 use std::fs;
252 use std::path::Path;
253
254 const HEADER: &str = "eyJhbGciOiJub25lIiwidHlwIjoiSldUIn0";
255 const PAYLOAD_ALPHA: &str = "eyJzdWIiOiJ1c2VyXzEyMyIsImVtYWlsIjoiYWxwaGFAZXhhbXBsZS5jb20iLCJodHRwczovL2FwaS5vcGVuYWkuY29tL2F1dGgiOnsiY2hhdGdwdF91c2VyX2lkIjoidXNlcl8xMjMiLCJlbWFpbCI6ImFscGhhQGV4YW1wbGUuY29tIn19";
256
257 fn token(payload: &str) -> String {
258 format!("{HEADER}.{payload}.sig")
259 }
260
261 fn auth_json(
262 payload: &str,
263 account_id: &str,
264 refresh_token: &str,
265 last_refresh: &str,
266 ) -> String {
267 format!(
268 r#"{{"tokens":{{"access_token":"{}","id_token":"{}","refresh_token":"{}","account_id":"{}"}},"last_refresh":"{}"}}"#,
269 token(payload),
270 token(payload),
271 refresh_token,
272 account_id,
273 last_refresh
274 )
275 }
276
277 fn set_cache_env(
278 lock: &GlobalStateLock,
279 secret_dir: &Path,
280 cache_root: &Path,
281 ) -> (EnvGuard, EnvGuard) {
282 let secret = EnvGuard::set(
283 lock,
284 "CODEX_SECRET_DIR",
285 secret_dir.to_str().expect("secret dir path"),
286 );
287 let cache = EnvGuard::set(
288 lock,
289 "ZSH_CACHE_DIR",
290 cache_root.to_str().expect("cache root path"),
291 );
292 (secret, cache)
293 }
294
295 #[test]
296 fn clear_starship_cache_rejects_relative_cache_root() {
297 let lock = GlobalStateLock::new();
298 let _cache = EnvGuard::set(&lock, "ZSH_CACHE_DIR", "relative/cache");
299
300 let err = clear_starship_cache().expect_err("relative cache root should fail");
301 assert!(err.to_string().contains("non-absolute cache root"));
302 }
303
304 #[test]
305 fn clear_starship_cache_rejects_root_cache_path() {
306 let lock = GlobalStateLock::new();
307 let _cache = EnvGuard::set(&lock, "ZSH_CACHE_DIR", "/");
308
309 let err = clear_starship_cache().expect_err("root cache path should fail");
310 assert!(err.to_string().contains("invalid cache root"));
311 }
312
313 #[test]
314 fn clear_starship_cache_removes_only_starship_cache_dir() {
315 let lock = GlobalStateLock::new();
316 let dir = tempfile::TempDir::new().expect("tempdir");
317 let cache_root = dir.path().join("cache-root");
318 let remove_dir = cache_root.join("codex").join("starship-rate-limits");
319 let keep_dir = cache_root.join("codex").join("secrets");
320 fs::create_dir_all(&remove_dir).expect("remove dir");
321 fs::create_dir_all(&keep_dir).expect("keep dir");
322 fs::write(
323 remove_dir.join("alpha.kv"),
324 "weekly_remaining=1\nweekly_reset_epoch=2",
325 )
326 .expect("write cached file");
327 fs::write(keep_dir.join("keep.txt"), "keep").expect("write keep file");
328 let _cache = EnvGuard::set(
329 &lock,
330 "ZSH_CACHE_DIR",
331 cache_root.to_str().expect("cache root path"),
332 );
333
334 clear_starship_cache().expect("clear cache");
335
336 assert!(!remove_dir.exists(), "starship cache dir should be removed");
337 assert!(keep_dir.is_dir(), "non-target cache dir should remain");
338 }
339
340 #[test]
341 fn cache_file_for_secret_target_uses_sanitized_secret_name() {
342 let lock = GlobalStateLock::new();
343 let dir = tempfile::TempDir::new().expect("tempdir");
344 let secret_dir = dir.path().join("secrets");
345 let cache_root = dir.path().join("cache");
346 fs::create_dir_all(&secret_dir).expect("secret dir");
347 fs::create_dir_all(&cache_root).expect("cache root");
348 let _env = set_cache_env(&lock, &secret_dir, &cache_root);
349
350 let target = secret_dir.join("My.Secret+Name.json");
351 fs::write(&target, "{}").expect("write secret file");
352
353 let cache_file = cache_file_for_target(&target).expect("cache file");
354 assert_eq!(
355 cache_file,
356 cache_root
357 .join("codex")
358 .join("starship-rate-limits")
359 .join("my_secret_name.kv")
360 );
361 }
362
363 #[test]
364 fn cache_file_for_non_secret_target_falls_back_to_hashed_key() {
365 let lock = GlobalStateLock::new();
366 let dir = tempfile::TempDir::new().expect("tempdir");
367 let secret_dir = dir.path().join("secrets");
368 let cache_root = dir.path().join("cache");
369 fs::create_dir_all(&secret_dir).expect("secret dir");
370 fs::create_dir_all(&cache_root).expect("cache root");
371 let _env = set_cache_env(&lock, &secret_dir, &cache_root);
372
373 let target = dir.path().join("auth.json");
374 fs::write(&target, "{\"tokens\":{\"access_token\":\"tok\"}}").expect("write auth file");
375
376 let hash = codex_fs::sha256_file(&target).expect("sha256");
377 let cache_file = cache_file_for_target(&target).expect("cache file");
378 assert_eq!(
379 cache_file,
380 cache_root
381 .join("codex")
382 .join("starship-rate-limits")
383 .join(format!("auth_{hash}.kv"))
384 );
385 }
386
387 #[test]
388 fn cache_file_for_auth_target_reuses_matching_secret_identity() {
389 let lock = GlobalStateLock::new();
390 let dir = tempfile::TempDir::new().expect("tempdir");
391 let secret_dir = dir.path().join("secrets");
392 let cache_root = dir.path().join("cache");
393 fs::create_dir_all(&secret_dir).expect("secret dir");
394 fs::create_dir_all(&cache_root).expect("cache root");
395 let _env = set_cache_env(&lock, &secret_dir, &cache_root);
396
397 let target = dir.path().join("auth.json");
398 let target_content = auth_json(
399 PAYLOAD_ALPHA,
400 "acct_001",
401 "refresh_auth",
402 "2025-01-20T12:34:56Z",
403 );
404 fs::write(&target, target_content).expect("write auth file");
405
406 let secret_file = secret_dir.join("Alpha Team.json");
407 let secret_content = auth_json(
408 PAYLOAD_ALPHA,
409 "acct_001",
410 "refresh_secret",
411 "2025-01-21T12:34:56Z",
412 );
413 fs::write(&secret_file, secret_content).expect("write matching secret file");
414
415 let cache_file = cache_file_for_target(&target).expect("cache file");
416 assert_eq!(
417 cache_file.file_name().and_then(|name| name.to_str()),
418 Some("alpha_team.kv")
419 );
420 assert_eq!(
421 secret_name_for_target(&target),
422 Some("Alpha Team".to_string())
423 );
424 }
425
426 #[test]
427 fn write_then_read_cache_entry_preserves_optional_non_weekly_reset_epoch() {
428 let lock = GlobalStateLock::new();
429 let dir = tempfile::TempDir::new().expect("tempdir");
430 let secret_dir = dir.path().join("secrets");
431 let cache_root = dir.path().join("cache");
432 fs::create_dir_all(&secret_dir).expect("secret dir");
433 fs::create_dir_all(&cache_root).expect("cache root");
434 let _env = set_cache_env(&lock, &secret_dir, &cache_root);
435
436 let target = secret_dir.join("alpha.json");
437 fs::write(&target, "{}").expect("write target");
438
439 write_starship_cache(
440 &target,
441 1700000000,
442 "5h",
443 91,
444 12,
445 1700600000,
446 Some(1700003600),
447 )
448 .expect("write cache");
449
450 let entry = read_cache_entry(&target).expect("read cache");
451 assert_eq!(entry.non_weekly_label, "5h");
452 assert_eq!(entry.non_weekly_remaining, 91);
453 assert_eq!(entry.non_weekly_reset_epoch, Some(1700003600));
454 assert_eq!(entry.weekly_remaining, 12);
455 assert_eq!(entry.weekly_reset_epoch, 1700600000);
456 }
457
458 #[test]
459 fn write_cache_omits_optional_non_weekly_reset_epoch_when_absent() {
460 let lock = GlobalStateLock::new();
461 let dir = tempfile::TempDir::new().expect("tempdir");
462 let secret_dir = dir.path().join("secrets");
463 let cache_root = dir.path().join("cache");
464 fs::create_dir_all(&secret_dir).expect("secret dir");
465 fs::create_dir_all(&cache_root).expect("cache root");
466 let _env = set_cache_env(&lock, &secret_dir, &cache_root);
467
468 let target = secret_dir.join("alpha.json");
469 fs::write(&target, "{}").expect("write target");
470
471 write_starship_cache(&target, 1700000000, "daily", 45, 9, 1700600000, None)
472 .expect("write cache");
473
474 let cache_file = cache_file_for_target(&target).expect("cache path");
475 let content = fs::read_to_string(&cache_file).expect("read cache file");
476 assert!(!content.contains("non_weekly_reset_epoch="));
477
478 let entry = read_cache_entry(&target).expect("read cache");
479 assert_eq!(entry.non_weekly_label, "daily");
480 assert_eq!(entry.non_weekly_reset_epoch, None);
481 }
482
483 #[test]
484 fn read_cache_entry_reports_missing_weekly_data() {
485 let lock = GlobalStateLock::new();
486 let dir = tempfile::TempDir::new().expect("tempdir");
487 let secret_dir = dir.path().join("secrets");
488 let cache_root = dir.path().join("cache");
489 fs::create_dir_all(&secret_dir).expect("secret dir");
490 fs::create_dir_all(&cache_root).expect("cache root");
491 let _env = set_cache_env(&lock, &secret_dir, &cache_root);
492
493 let target = secret_dir.join("alpha.json");
494 fs::write(&target, "{}").expect("write target");
495 let cache_file = cache_file_for_target(&target).expect("cache path");
496 fs::create_dir_all(cache_file.parent().expect("cache parent")).expect("cache parent dir");
497 fs::write(
498 &cache_file,
499 "fetched_at=1\nnon_weekly_label=5h\nnon_weekly_remaining=90\nweekly_remaining=1\n",
500 )
501 .expect("write invalid cache");
502
503 let err = read_cache_entry(&target)
504 .err()
505 .expect("missing weekly reset should fail");
506 assert!(err.to_string().contains("missing weekly data"));
507 }
508
509 #[test]
510 fn read_cache_entry_reports_missing_non_weekly_data() {
511 let lock = GlobalStateLock::new();
512 let dir = tempfile::TempDir::new().expect("tempdir");
513 let secret_dir = dir.path().join("secrets");
514 let cache_root = dir.path().join("cache");
515 fs::create_dir_all(&secret_dir).expect("secret dir");
516 fs::create_dir_all(&cache_root).expect("cache root");
517 let _env = set_cache_env(&lock, &secret_dir, &cache_root);
518
519 let target = secret_dir.join("alpha.json");
520 fs::write(&target, "{}").expect("write target");
521 let cache_file = cache_file_for_target(&target).expect("cache path");
522 fs::create_dir_all(cache_file.parent().expect("cache parent")).expect("cache parent dir");
523 fs::write(
524 &cache_file,
525 "fetched_at=1\nweekly_remaining=1\nweekly_reset_epoch=1700600000\n",
526 )
527 .expect("write invalid cache");
528
529 let err = read_cache_entry(&target)
530 .err()
531 .expect("missing non-weekly fields should fail");
532 assert!(err.to_string().contains("missing non-weekly data"));
533 }
534}