1use std::fs;
2use std::path::{Path, PathBuf};
3
4use crate::auth;
5use crate::fs as gemini_fs;
6use crate::paths;
7use crate::rate_limits::client::{UsageRequest, fetch_usage};
8
9pub mod ansi;
10pub mod client;
11pub mod render;
12
13#[derive(Clone, Debug, Default)]
14pub struct RateLimitsOptions {
15 pub clear_cache: bool,
16 pub debug: bool,
17 pub cached: bool,
18 pub no_refresh_auth: bool,
19 pub json: bool,
20 pub one_line: bool,
21 pub all: bool,
22 pub async_mode: bool,
23 pub jobs: Option<String>,
24 pub secret: Option<String>,
25}
26
27#[derive(Clone, Debug)]
28pub struct CacheEntry {
29 pub non_weekly_label: String,
30 pub non_weekly_remaining: i64,
31 pub non_weekly_reset_epoch: Option<i64>,
32 pub weekly_remaining: i64,
33 pub weekly_reset_epoch: i64,
34}
35
36pub const DIAG_SCHEMA_VERSION: &str = "gemini-cli.diag.rate-limits.v1";
37pub const DIAG_COMMAND: &str = "diag rate-limits";
38
39#[derive(Clone, Debug)]
40struct RateLimitSummary {
41 non_weekly_label: String,
42 non_weekly_remaining: i64,
43 non_weekly_reset_epoch: Option<i64>,
44 weekly_remaining: i64,
45 weekly_reset_epoch: i64,
46}
47
48#[derive(Clone, Debug)]
49struct JsonResultItem {
50 name: String,
51 target_file: String,
52 status: String,
53 ok: bool,
54 source: String,
55 summary: Option<RateLimitSummary>,
56 raw_usage: Option<String>,
57 error_code: Option<String>,
58 error_message: Option<String>,
59}
60
61struct Row {
62 name: String,
63 window_label: String,
64 non_weekly_remaining: i64,
65 non_weekly_reset_epoch: Option<i64>,
66 weekly_remaining: i64,
67 weekly_reset_epoch: Option<i64>,
68}
69
70impl Row {
71 fn empty(name: String) -> Self {
72 Self {
73 name,
74 window_label: String::new(),
75 non_weekly_remaining: -1,
76 non_weekly_reset_epoch: None,
77 weekly_remaining: -1,
78 weekly_reset_epoch: None,
79 }
80 }
81
82 fn sort_key(&self) -> (i32, i64, String) {
83 if let Some(epoch) = self.weekly_reset_epoch {
84 (0, epoch, self.name.clone())
85 } else {
86 (1, i64::MAX, self.name.clone())
87 }
88 }
89}
90
91pub fn run(args: &RateLimitsOptions) -> i32 {
92 let cached_mode = args.cached;
93 let output_json = args.json;
94
95 if args.async_mode {
96 if !cached_mode {
97 maybe_sync_all_mode_auth_silent(args.debug);
98 }
99 if output_json {
100 return run_async_json_mode(args);
101 }
102 return run_async_mode(args);
103 }
104
105 if cached_mode {
106 if output_json {
107 emit_error_json(
108 "invalid-flag-combination",
109 "gemini-rate-limits: --json is not supported with --cached",
110 Some(json_obj(vec![(
111 "flags".to_string(),
112 json_array(vec![json_string("--json"), json_string("--cached")]),
113 )])),
114 );
115 return 64;
116 }
117 if args.clear_cache {
118 eprintln!("gemini-rate-limits: -c is not compatible with --cached");
119 return 64;
120 }
121 }
122
123 if output_json && args.one_line {
124 emit_error_json(
125 "invalid-flag-combination",
126 "gemini-rate-limits: --one-line is not compatible with --json",
127 Some(json_obj(vec![(
128 "flags".to_string(),
129 json_array(vec![json_string("--one-line"), json_string("--json")]),
130 )])),
131 );
132 return 64;
133 }
134
135 if args.clear_cache
136 && let Err(err) = clear_starship_cache()
137 {
138 if output_json {
139 emit_error_json("cache-clear-failed", &err, None);
140 } else {
141 eprintln!("{err}");
142 }
143 return 1;
144 }
145
146 let default_all_enabled = env_truthy("GEMINI_RATE_LIMITS_DEFAULT_ALL_ENABLED");
147 let all_mode = args.all
148 || (!args.cached
149 && !output_json
150 && args.secret.is_none()
151 && default_all_enabled
152 && !args.async_mode);
153
154 if all_mode {
155 if !cached_mode {
156 maybe_sync_all_mode_auth_silent(args.debug);
157 }
158 if args.secret.is_some() {
159 eprintln!(
160 "gemini-rate-limits: usage: gemini-rate-limits [-c] [--cached] [--no-refresh-auth] [--json] [--one-line] [--all] [secret.json]"
161 );
162 return 64;
163 }
164 if output_json {
165 return run_all_json_mode(args, cached_mode);
166 }
167 return run_all_mode(args, cached_mode);
168 }
169
170 run_single_mode(args, cached_mode, output_json)
171}
172
173fn maybe_sync_all_mode_auth_silent(debug_mode: bool) {
174 if let Err(err) = sync_auth_silent()
175 && debug_mode
176 {
177 eprintln!("{err}");
178 }
179}
180
181fn sync_auth_silent() -> Result<(), String> {
182 let auth_file = match paths::resolve_auth_file() {
183 Some(path) => path,
184 None => return Ok(()),
185 };
186
187 if !auth_file.is_file() {
188 return Ok(());
189 }
190
191 let auth_key = match auth::identity_key_from_auth_file(&auth_file) {
192 Ok(Some(key)) => key,
193 _ => return Ok(()),
194 };
195
196 let auth_last_refresh = auth::last_refresh_from_auth_file(&auth_file).ok().flatten();
197 let auth_contents = fs::read(&auth_file)
198 .map_err(|_| format!("gemini-rate-limits: failed to read {}", auth_file.display()))?;
199
200 if let Some(secret_dir) = paths::resolve_secret_dir()
201 && let Ok(entries) = fs::read_dir(&secret_dir)
202 {
203 for entry in entries.flatten() {
204 let path = entry.path();
205 if path.extension().and_then(|value| value.to_str()) != Some("json") {
206 continue;
207 }
208
209 let candidate_key = match auth::identity_key_from_auth_file(&path) {
210 Ok(Some(key)) => key,
211 _ => continue,
212 };
213 if candidate_key != auth_key {
214 continue;
215 }
216
217 let secret_contents = fs::read(&path)
218 .map_err(|_| format!("gemini-rate-limits: failed to read {}", path.display()))?;
219 if secret_contents == auth_contents {
220 continue;
221 }
222
223 auth::write_atomic(&path, &auth_contents, auth::SECRET_FILE_MODE)
224 .map_err(|_| format!("gemini-rate-limits: failed to write {}", path.display()))?;
225
226 if let Some(timestamp_path) = sync_timestamp_path(&path) {
227 let _ = auth::write_timestamp(×tamp_path, auth_last_refresh.as_deref());
228 }
229 }
230 }
231
232 if let Some(auth_timestamp) = sync_timestamp_path(&auth_file) {
233 let _ = auth::write_timestamp(&auth_timestamp, auth_last_refresh.as_deref());
234 }
235
236 Ok(())
237}
238
239fn sync_timestamp_path(target_file: &Path) -> Option<PathBuf> {
240 let cache_dir = paths::resolve_secret_cache_dir()?;
241 let name = target_file
242 .file_name()
243 .and_then(|name| name.to_str())
244 .unwrap_or("auth.json");
245 Some(cache_dir.join(format!("{name}.timestamp")))
246}
247
248fn run_single_mode(args: &RateLimitsOptions, cached_mode: bool, output_json: bool) -> i32 {
249 let target_file = match resolve_single_target(args.secret.as_deref()) {
250 Ok(path) => path,
251 Err(message) => {
252 if output_json {
253 emit_error_json("target-not-found", &message, None);
254 } else {
255 eprintln!("{message}");
256 }
257 return 1;
258 }
259 };
260
261 if cached_mode {
262 let cache_entry = match read_cache_entry(&target_file) {
263 Ok(entry) => entry,
264 Err(err) => {
265 eprintln!("{err}");
266 return 1;
267 }
268 };
269 let name = secret_name_for_target(&target_file).unwrap_or_else(|| {
270 target_file
271 .file_stem()
272 .and_then(|value| value.to_str())
273 .unwrap_or("auth")
274 .to_string()
275 });
276 let summary = RateLimitSummary {
277 non_weekly_label: cache_entry.non_weekly_label,
278 non_weekly_remaining: cache_entry.non_weekly_remaining,
279 non_weekly_reset_epoch: cache_entry.non_weekly_reset_epoch,
280 weekly_remaining: cache_entry.weekly_remaining,
281 weekly_reset_epoch: cache_entry.weekly_reset_epoch,
282 };
283 if args.one_line {
284 let line = render_line_for_summary(&name, &summary, true, "%m-%d %H:%M");
285 println!("{line}");
286 } else {
287 print_rate_limits_remaining(&summary, "%m-%d %H:%M");
288 }
289 return 0;
290 }
291
292 match collect_summary_from_network(&target_file, !args.no_refresh_auth) {
293 Ok((summary, raw_usage)) => {
294 if output_json {
295 let item = JsonResultItem {
296 name: secret_name_for_target(&target_file)
297 .unwrap_or_else(|| "auth".to_string()),
298 target_file: target_file_name(&target_file),
299 status: "ok".to_string(),
300 ok: true,
301 source: "network".to_string(),
302 summary: Some(summary.clone()),
303 raw_usage,
304 error_code: None,
305 error_message: None,
306 };
307 emit_single_envelope("single", true, &item);
308 } else {
309 let name = secret_name_for_target(&target_file).unwrap_or_else(|| {
310 target_file
311 .file_stem()
312 .and_then(|value| value.to_str())
313 .unwrap_or("auth")
314 .to_string()
315 });
316 if args.one_line {
317 let line =
318 render_line_for_summary(&name, &summary, args.one_line, "%m-%d %H:%M");
319 println!("{line}");
320 } else {
321 print_rate_limits_remaining(&summary, "%m-%d %H:%M");
322 }
323 }
324 0
325 }
326 Err(err) => {
327 if output_json {
328 emit_error_json(&err.code, &err.message, err.details);
329 } else {
330 eprintln!("{}", err.message);
331 }
332 err.exit_code
333 }
334 }
335}
336
337fn run_all_mode(args: &RateLimitsOptions, cached_mode: bool) -> i32 {
338 let secret_files = match collect_secret_files() {
339 Ok(value) => value,
340 Err(err) => {
341 eprintln!("{err}");
342 return 1;
343 }
344 };
345
346 let current_name = current_secret_basename(&secret_files);
347
348 let mut rc = 0;
349 let mut rows: Vec<Row> = Vec::new();
350 let mut window_labels = std::collections::HashSet::new();
351
352 for target in secret_files {
353 let name = secret_name_for_target(&target).unwrap_or_else(|| target_file_name(&target));
354 let mut row = Row::empty(name.clone());
355
356 if cached_mode {
357 match read_cache_entry(&target) {
358 Ok(summary) => {
359 row.window_label = summary.non_weekly_label.clone();
360 row.non_weekly_remaining = summary.non_weekly_remaining;
361 row.non_weekly_reset_epoch = summary.non_weekly_reset_epoch;
362 row.weekly_remaining = summary.weekly_remaining;
363 row.weekly_reset_epoch = Some(summary.weekly_reset_epoch);
364 window_labels.insert(row.window_label.clone());
365 }
366 Err(err) => {
367 eprintln!("{name}: {err}");
368 rc = 1;
369 }
370 }
371 rows.push(row);
372 continue;
373 }
374
375 match collect_summary_from_network(&target, !args.no_refresh_auth) {
376 Ok((summary, _raw)) => {
377 row.window_label = summary.non_weekly_label.clone();
378 row.non_weekly_remaining = summary.non_weekly_remaining;
379 row.non_weekly_reset_epoch = summary.non_weekly_reset_epoch;
380 row.weekly_remaining = summary.weekly_remaining;
381 row.weekly_reset_epoch = Some(summary.weekly_reset_epoch);
382 window_labels.insert(row.window_label.clone());
383 }
384 Err(err) => {
385 eprintln!("{name}: {}", err.message);
386 rc = 1;
387 }
388 }
389 rows.push(row);
390 }
391
392 println!("\n🚦 Gemini rate limits for all accounts\n");
393
394 let mut non_weekly_header = "Non-weekly".to_string();
395 let multiple_labels = window_labels.len() != 1;
396 if !multiple_labels && let Some(label) = window_labels.iter().next() {
397 non_weekly_header = label.clone();
398 }
399
400 let now_epoch = now_epoch_seconds();
401
402 println!(
403 "{:<15} {:>8} {:>7} {:>8} {:>7} {:<18}",
404 "Name", non_weekly_header, "Left", "Weekly", "Left", "Reset"
405 );
406 println!("----------------------------------------------------------------------------");
407
408 rows.sort_by_key(|row| row.sort_key());
409
410 for row in rows {
411 let display_non_weekly = if multiple_labels && !row.window_label.is_empty() {
412 if row.non_weekly_remaining >= 0 {
413 format!("{}:{}%", row.window_label, row.non_weekly_remaining)
414 } else {
415 "-".to_string()
416 }
417 } else if row.non_weekly_remaining >= 0 {
418 format!("{}%", row.non_weekly_remaining)
419 } else {
420 "-".to_string()
421 };
422
423 let non_weekly_left = row
424 .non_weekly_reset_epoch
425 .and_then(|epoch| render::format_until_epoch_compact(epoch, now_epoch))
426 .unwrap_or_else(|| "-".to_string());
427 let weekly_left = row
428 .weekly_reset_epoch
429 .and_then(|epoch| render::format_until_epoch_compact(epoch, now_epoch))
430 .unwrap_or_else(|| "-".to_string());
431 let reset_display = row
432 .weekly_reset_epoch
433 .and_then(render::format_epoch_local_datetime_with_offset)
434 .unwrap_or_else(|| "-".to_string());
435
436 let non_weekly_display = ansi::format_percent_cell(&display_non_weekly, 8, None);
437 let weekly_display = if row.weekly_remaining >= 0 {
438 ansi::format_percent_cell(&format!("{}%", row.weekly_remaining), 8, None)
439 } else {
440 ansi::format_percent_cell("-", 8, None)
441 };
442
443 let is_current = current_name.as_deref() == Some(row.name.as_str());
444 let name_display = ansi::format_name_cell(&row.name, 15, is_current, None);
445
446 println!(
447 "{} {} {:>7} {} {:>7} {:<18}",
448 name_display,
449 non_weekly_display,
450 non_weekly_left,
451 weekly_display,
452 weekly_left,
453 reset_display
454 );
455 }
456
457 rc
458}
459
460fn run_all_json_mode(args: &RateLimitsOptions, cached_mode: bool) -> i32 {
461 let secret_files = match collect_secret_files() {
462 Ok(value) => value,
463 Err(err) => {
464 emit_error_json("secret-discovery-failed", &err, None);
465 return 1;
466 }
467 };
468
469 let mut items: Vec<JsonResultItem> = Vec::new();
470 let mut rc = 0;
471
472 for target in secret_files {
473 let name = secret_name_for_target(&target).unwrap_or_else(|| target_file_name(&target));
474 if cached_mode {
475 match read_cache_entry(&target) {
476 Ok(entry) => items.push(JsonResultItem {
477 name,
478 target_file: target_file_name(&target),
479 status: "ok".to_string(),
480 ok: true,
481 source: "cache".to_string(),
482 summary: Some(RateLimitSummary {
483 non_weekly_label: entry.non_weekly_label,
484 non_weekly_remaining: entry.non_weekly_remaining,
485 non_weekly_reset_epoch: entry.non_weekly_reset_epoch,
486 weekly_remaining: entry.weekly_remaining,
487 weekly_reset_epoch: entry.weekly_reset_epoch,
488 }),
489 raw_usage: None,
490 error_code: None,
491 error_message: None,
492 }),
493 Err(err) => {
494 rc = 1;
495 items.push(JsonResultItem {
496 name,
497 target_file: target_file_name(&target),
498 status: "error".to_string(),
499 ok: false,
500 source: "cache".to_string(),
501 summary: None,
502 raw_usage: None,
503 error_code: Some("cache-read-failed".to_string()),
504 error_message: Some(err),
505 });
506 }
507 }
508 continue;
509 }
510
511 match collect_summary_from_network(&target, !args.no_refresh_auth) {
512 Ok((summary, raw_usage)) => items.push(JsonResultItem {
513 name,
514 target_file: target_file_name(&target),
515 status: "ok".to_string(),
516 ok: true,
517 source: "network".to_string(),
518 summary: Some(summary),
519 raw_usage,
520 error_code: None,
521 error_message: None,
522 }),
523 Err(err) => {
524 rc = 1;
525 items.push(JsonResultItem {
526 name,
527 target_file: target_file_name(&target),
528 status: "error".to_string(),
529 ok: false,
530 source: "network".to_string(),
531 summary: None,
532 raw_usage: None,
533 error_code: Some(err.code),
534 error_message: Some(err.message),
535 });
536 }
537 }
538 }
539
540 items.sort_by(|a, b| a.name.cmp(&b.name));
541 emit_collection_envelope("all", rc == 0, &items);
542 rc
543}
544
545fn run_async_mode(args: &RateLimitsOptions) -> i32 {
546 if args.one_line {
547 eprintln!("gemini-rate-limits: --async does not support --one-line");
548 return 64;
549 }
550 if let Some(secret) = args.secret.as_deref() {
551 eprintln!("gemini-rate-limits: --async does not accept positional args: {secret}");
552 eprintln!("hint: async always queries all secrets under GEMINI_SECRET_DIR");
553 return 64;
554 }
555 if args.clear_cache && args.cached {
556 eprintln!("gemini-rate-limits: --async: -c is not compatible with --cached");
557 return 64;
558 }
559 if args.clear_cache
560 && let Err(err) = clear_starship_cache()
561 {
562 eprintln!("{err}");
563 return 1;
564 }
565 run_all_mode(args, args.cached)
566}
567
568fn run_async_json_mode(args: &RateLimitsOptions) -> i32 {
569 if args.one_line {
570 emit_error_json(
571 "invalid-flag-combination",
572 "gemini-rate-limits: --async does not support --one-line",
573 Some(json_obj(vec![
574 ("flag".to_string(), json_string("--one-line")),
575 ("mode".to_string(), json_string("async")),
576 ])),
577 );
578 return 64;
579 }
580 if let Some(secret) = args.secret.as_deref() {
581 emit_error_json(
582 "invalid-positional-arg",
583 &format!("gemini-rate-limits: --async does not accept positional args: {secret}"),
584 Some(json_obj(vec![
585 ("secret".to_string(), json_string(secret)),
586 ("mode".to_string(), json_string("async")),
587 ])),
588 );
589 return 64;
590 }
591 if args.clear_cache && args.cached {
592 emit_error_json(
593 "invalid-flag-combination",
594 "gemini-rate-limits: --async: -c is not compatible with --cached",
595 Some(json_obj(vec![(
596 "flags".to_string(),
597 json_array(vec![
598 json_string("--async"),
599 json_string("--cached"),
600 json_string("-c"),
601 ]),
602 )])),
603 );
604 return 64;
605 }
606 if args.clear_cache
607 && let Err(err) = clear_starship_cache()
608 {
609 emit_error_json("cache-clear-failed", &err, None);
610 return 1;
611 }
612
613 let secret_files = match collect_secret_files() {
614 Ok(value) => value,
615 Err(err) => {
616 emit_error_json("secret-discovery-failed", &err, None);
617 return 1;
618 }
619 };
620
621 let mut items: Vec<JsonResultItem> = Vec::new();
622 let mut rc = 0;
623 for target in secret_files {
624 let name = secret_name_for_target(&target).unwrap_or_else(|| target_file_name(&target));
625 match collect_summary_from_network(&target, !args.no_refresh_auth) {
626 Ok((summary, raw_usage)) => items.push(JsonResultItem {
627 name,
628 target_file: target_file_name(&target),
629 status: "ok".to_string(),
630 ok: true,
631 source: "network".to_string(),
632 summary: Some(summary),
633 raw_usage,
634 error_code: None,
635 error_message: None,
636 }),
637 Err(err) => {
638 if err.code == "missing-access-token"
639 && let Ok(cached) = read_cache_entry(&target)
640 {
641 items.push(JsonResultItem {
642 name,
643 target_file: target_file_name(&target),
644 status: "ok".to_string(),
645 ok: true,
646 source: "cache-fallback".to_string(),
647 summary: Some(RateLimitSummary {
648 non_weekly_label: cached.non_weekly_label,
649 non_weekly_remaining: cached.non_weekly_remaining,
650 non_weekly_reset_epoch: cached.non_weekly_reset_epoch,
651 weekly_remaining: cached.weekly_remaining,
652 weekly_reset_epoch: cached.weekly_reset_epoch,
653 }),
654 raw_usage: None,
655 error_code: None,
656 error_message: None,
657 });
658 continue;
659 }
660
661 rc = 1;
662 items.push(JsonResultItem {
663 name,
664 target_file: target_file_name(&target),
665 status: "error".to_string(),
666 ok: false,
667 source: "network".to_string(),
668 summary: None,
669 raw_usage: None,
670 error_code: Some(err.code),
671 error_message: Some(err.message),
672 });
673 }
674 }
675 }
676 items.sort_by(|a, b| a.name.cmp(&b.name));
677 emit_collection_envelope("async", rc == 0, &items);
678 rc
679}
680
681pub fn clear_starship_cache() -> Result<(), String> {
682 let root =
683 cache_root().ok_or_else(|| "gemini-rate-limits: cache root unavailable".to_string())?;
684 if !root.is_absolute() {
685 return Err(format!(
686 "gemini-rate-limits: refusing to clear cache with non-absolute cache root: {}",
687 root.display()
688 ));
689 }
690 if root == Path::new("/") {
691 return Err(format!(
692 "gemini-rate-limits: refusing to clear cache with invalid cache root: {}",
693 root.display()
694 ));
695 }
696
697 let cache_dir = root.join("gemini").join("starship-rate-limits");
698 let cache_dir_str = cache_dir.to_string_lossy();
699 if !cache_dir_str.ends_with("/gemini/starship-rate-limits") {
700 return Err(format!(
701 "gemini-rate-limits: refusing to clear unexpected cache dir: {}",
702 cache_dir.display()
703 ));
704 }
705
706 if cache_dir.is_dir() {
707 let _ = fs::remove_dir_all(&cache_dir);
708 }
709 Ok(())
710}
711
712pub fn cache_file_for_target(target_file: &Path) -> Result<PathBuf, String> {
713 let cache_dir = starship_cache_dir()
714 .ok_or_else(|| "gemini-rate-limits: cache dir unavailable".to_string())?;
715
716 if let Some(secret_dir) = paths::resolve_secret_dir() {
717 if target_file.starts_with(&secret_dir) {
718 let display = secret_file_basename(target_file)?;
719 let key = cache_key(&display)?;
720 return Ok(cache_dir.join(format!("{key}.kv")));
721 }
722
723 if let Some(secret_name) = secret_name_for_auth(target_file, &secret_dir) {
724 let key = cache_key(&secret_name)?;
725 return Ok(cache_dir.join(format!("{key}.kv")));
726 }
727 }
728
729 let hash = gemini_fs::sha256_file(target_file).map_err(|err| err.to_string())?;
730 Ok(cache_dir.join(format!("auth_{}.kv", hash.to_lowercase())))
731}
732
733pub fn secret_name_for_target(target_file: &Path) -> Option<String> {
734 let secret_dir = paths::resolve_secret_dir()?;
735 if target_file.starts_with(&secret_dir) {
736 return secret_file_basename(target_file).ok();
737 }
738 secret_name_for_auth(target_file, &secret_dir)
739}
740
741pub fn read_cache_entry(target_file: &Path) -> Result<CacheEntry, String> {
742 let cache_file = cache_file_for_target(target_file)?;
743 if !cache_file.is_file() {
744 return Err(format!(
745 "gemini-rate-limits: cache not found (run gemini-rate-limits without --cached, or gemini-cli starship, to populate): {}",
746 cache_file.display()
747 ));
748 }
749
750 let content = fs::read_to_string(&cache_file).map_err(|_| {
751 format!(
752 "gemini-rate-limits: failed to read cache: {}",
753 cache_file.display()
754 )
755 })?;
756 let mut non_weekly_label: Option<String> = None;
757 let mut non_weekly_remaining: Option<i64> = None;
758 let mut non_weekly_reset_epoch: Option<i64> = None;
759 let mut weekly_remaining: Option<i64> = None;
760 let mut weekly_reset_epoch: Option<i64> = None;
761
762 for line in content.lines() {
763 if let Some(value) = line.strip_prefix("non_weekly_label=") {
764 non_weekly_label = Some(value.to_string());
765 } else if let Some(value) = line.strip_prefix("non_weekly_remaining=") {
766 non_weekly_remaining = value.parse::<i64>().ok();
767 } else if let Some(value) = line.strip_prefix("non_weekly_reset_epoch=") {
768 non_weekly_reset_epoch = value.parse::<i64>().ok();
769 } else if let Some(value) = line.strip_prefix("weekly_remaining=") {
770 weekly_remaining = value.parse::<i64>().ok();
771 } else if let Some(value) = line.strip_prefix("weekly_reset_epoch=") {
772 weekly_reset_epoch = value.parse::<i64>().ok();
773 }
774 }
775
776 let non_weekly_label = match non_weekly_label {
777 Some(value) if !value.trim().is_empty() => value,
778 _ => {
779 return Err(format!(
780 "gemini-rate-limits: invalid cache (missing non-weekly data): {}",
781 cache_file.display()
782 ));
783 }
784 };
785 let non_weekly_remaining = match non_weekly_remaining {
786 Some(value) => value,
787 None => {
788 return Err(format!(
789 "gemini-rate-limits: invalid cache (missing non-weekly data): {}",
790 cache_file.display()
791 ));
792 }
793 };
794 let weekly_remaining = match weekly_remaining {
795 Some(value) => value,
796 None => {
797 return Err(format!(
798 "gemini-rate-limits: invalid cache (missing weekly data): {}",
799 cache_file.display()
800 ));
801 }
802 };
803 let weekly_reset_epoch = match weekly_reset_epoch {
804 Some(value) => value,
805 None => {
806 return Err(format!(
807 "gemini-rate-limits: invalid cache (missing weekly data): {}",
808 cache_file.display()
809 ));
810 }
811 };
812
813 Ok(CacheEntry {
814 non_weekly_label,
815 non_weekly_remaining,
816 non_weekly_reset_epoch,
817 weekly_remaining,
818 weekly_reset_epoch,
819 })
820}
821
822pub fn write_starship_cache(
823 target_file: &Path,
824 fetched_at_epoch: i64,
825 non_weekly_label: &str,
826 non_weekly_remaining: i64,
827 weekly_remaining: i64,
828 weekly_reset_epoch: i64,
829 non_weekly_reset_epoch: Option<i64>,
830) -> Result<(), String> {
831 let cache_file = cache_file_for_target(target_file)?;
832 if let Some(parent) = cache_file.parent() {
833 fs::create_dir_all(parent).map_err(|err| err.to_string())?;
834 }
835
836 let mut lines = Vec::new();
837 lines.push(format!("fetched_at={fetched_at_epoch}"));
838 lines.push(format!("non_weekly_label={non_weekly_label}"));
839 lines.push(format!("non_weekly_remaining={non_weekly_remaining}"));
840 if let Some(epoch) = non_weekly_reset_epoch {
841 lines.push(format!("non_weekly_reset_epoch={epoch}"));
842 }
843 lines.push(format!("weekly_remaining={weekly_remaining}"));
844 lines.push(format!("weekly_reset_epoch={weekly_reset_epoch}"));
845
846 let data = lines.join("\n");
847 gemini_fs::write_atomic(&cache_file, data.as_bytes(), gemini_fs::SECRET_FILE_MODE)
848 .map_err(|err| err.to_string())
849}
850
851const DEFAULT_CODE_ASSIST_ENDPOINT: &str = "https://cloudcode-pa.googleapis.com";
852const DEFAULT_CODE_ASSIST_API_VERSION: &str = "v1internal";
853const DEFAULT_CODE_ASSIST_PROJECT: &str = "projects/default";
854
855fn run_code_assist_endpoint() -> String {
856 env_non_empty("CODE_ASSIST_ENDPOINT")
857 .or_else(|| env_non_empty("GEMINI_CODE_ASSIST_ENDPOINT"))
858 .unwrap_or_else(|| DEFAULT_CODE_ASSIST_ENDPOINT.to_string())
859}
860
861fn run_code_assist_api_version() -> String {
862 env_non_empty("CODE_ASSIST_API_VERSION")
863 .or_else(|| env_non_empty("GEMINI_CODE_ASSIST_API_VERSION"))
864 .unwrap_or_else(|| DEFAULT_CODE_ASSIST_API_VERSION.to_string())
865}
866
867fn run_code_assist_project() -> String {
868 let raw = env_non_empty("GEMINI_CODE_ASSIST_PROJECT")
869 .or_else(|| env_non_empty("GOOGLE_CLOUD_PROJECT"))
870 .or_else(|| env_non_empty("GOOGLE_CLOUD_PROJECT_ID"))
871 .unwrap_or_else(|| DEFAULT_CODE_ASSIST_PROJECT.to_string());
872
873 if raw.starts_with("projects/") {
874 raw
875 } else {
876 format!("projects/{raw}")
877 }
878}
879
880fn run_connect_timeout() -> u64 {
881 std::env::var("GEMINI_RATE_LIMITS_CURL_CONNECT_TIMEOUT_SECONDS")
882 .ok()
883 .and_then(|raw| raw.trim().parse::<u64>().ok())
884 .unwrap_or(2)
885}
886
887fn run_max_time() -> u64 {
888 std::env::var("GEMINI_RATE_LIMITS_CURL_MAX_TIME_SECONDS")
889 .ok()
890 .and_then(|raw| raw.trim().parse::<u64>().ok())
891 .unwrap_or(8)
892}
893
894fn collect_summary_from_network(
895 target_file: &Path,
896 refresh_on_401: bool,
897) -> Result<(RateLimitSummary, Option<String>), RunError> {
898 let request = UsageRequest {
899 target_file: target_file.to_path_buf(),
900 refresh_on_401,
901 endpoint: run_code_assist_endpoint(),
902 api_version: run_code_assist_api_version(),
903 project: run_code_assist_project(),
904 connect_timeout_seconds: run_connect_timeout(),
905 max_time_seconds: run_max_time(),
906 };
907 let usage = fetch_usage(&request).map_err(|message| {
908 let (code, exit_code) = if message.contains("missing access_token") {
909 ("missing-access-token".to_string(), 2)
910 } else {
911 ("request-failed".to_string(), 3)
912 };
913 RunError {
914 code,
915 message,
916 details: None,
917 exit_code,
918 }
919 })?;
920
921 let usage_data = render::parse_usage(&usage.body).ok_or_else(|| RunError {
922 code: "invalid-usage-payload".to_string(),
923 message: "gemini-rate-limits: invalid usage payload".to_string(),
924 details: Some(json_obj(vec![(
925 "raw_usage".to_string(),
926 usage.body.clone(),
927 )])),
928 exit_code: 3,
929 })?;
930 let values = render::render_values(&usage_data);
931 let weekly = render::weekly_values(&values);
932 let summary = RateLimitSummary {
933 non_weekly_label: weekly.non_weekly_label.clone(),
934 non_weekly_remaining: weekly.non_weekly_remaining,
935 non_weekly_reset_epoch: weekly.non_weekly_reset_epoch,
936 weekly_remaining: weekly.weekly_remaining,
937 weekly_reset_epoch: weekly.weekly_reset_epoch,
938 };
939
940 let now_epoch = now_epoch_seconds();
941 if now_epoch > 0 {
942 let _ = write_starship_cache(
943 target_file,
944 now_epoch,
945 &summary.non_weekly_label,
946 summary.non_weekly_remaining,
947 summary.weekly_remaining,
948 summary.weekly_reset_epoch,
949 summary.non_weekly_reset_epoch,
950 );
951 }
952
953 let raw_usage = if usage.body.trim_start().starts_with('{') {
954 Some(usage.body)
955 } else {
956 None
957 };
958
959 Ok((summary, raw_usage))
960}
961
962fn collect_secret_files() -> Result<Vec<PathBuf>, String> {
963 let secret_dir = paths::resolve_secret_dir().unwrap_or_default();
964 if !secret_dir.is_dir() {
965 return Err(format!(
966 "gemini-rate-limits: GEMINI_SECRET_DIR not found: {}",
967 secret_dir.display()
968 ));
969 }
970
971 let mut files: Vec<PathBuf> = fs::read_dir(&secret_dir)
972 .map_err(|err| format!("gemini-rate-limits: failed to read GEMINI_SECRET_DIR: {err}"))?
973 .flatten()
974 .map(|entry| entry.path())
975 .filter(|path| path.extension().and_then(|value| value.to_str()) == Some("json"))
976 .collect();
977
978 files.sort();
979
980 if files.is_empty() {
981 return Err(format!(
982 "gemini-rate-limits: no secrets found under GEMINI_SECRET_DIR: {}",
983 secret_dir.display()
984 ));
985 }
986
987 Ok(files)
988}
989
990fn resolve_single_target(secret: Option<&str>) -> Result<PathBuf, String> {
991 if let Some(raw) = secret {
992 if raw.trim().is_empty() {
993 return Err("gemini-rate-limits: empty secret target".to_string());
994 }
995 let path = if raw.contains('/') || raw.starts_with('.') {
996 PathBuf::from(raw)
997 } else if let Some(secret_dir) = paths::resolve_secret_dir() {
998 let mut file = raw.to_string();
999 if !file.ends_with(".json") {
1000 file.push_str(".json");
1001 }
1002 secret_dir.join(file)
1003 } else {
1004 PathBuf::from(raw)
1005 };
1006
1007 if !path.is_file() {
1008 return Err(format!(
1009 "gemini-rate-limits: target file not found: {}",
1010 path.display()
1011 ));
1012 }
1013 return Ok(path);
1014 }
1015
1016 let auth = paths::resolve_auth_file().ok_or_else(|| {
1017 "gemini-rate-limits: GEMINI_AUTH_FILE is not configured and no secret provided".to_string()
1018 })?;
1019 if !auth.is_file() {
1020 return Err(format!(
1021 "gemini-rate-limits: target file not found: {}",
1022 auth.display()
1023 ));
1024 }
1025 Ok(auth)
1026}
1027
1028fn secret_file_basename(path: &Path) -> Result<String, String> {
1029 let file = path
1030 .file_name()
1031 .and_then(|name| name.to_str())
1032 .unwrap_or_default();
1033 let base = file.trim_end_matches(".json");
1034 if base.is_empty() {
1035 return Err("missing secret basename".to_string());
1036 }
1037 Ok(base.to_string())
1038}
1039
1040fn cache_key(name: &str) -> Result<String, String> {
1041 if name.is_empty() {
1042 return Err("missing cache key name".to_string());
1043 }
1044 let mut key = String::new();
1045 for ch in name.to_lowercase().chars() {
1046 if ch.is_ascii_alphanumeric() {
1047 key.push(ch);
1048 } else {
1049 key.push('_');
1050 }
1051 }
1052 while key.starts_with('_') {
1053 key.remove(0);
1054 }
1055 while key.ends_with('_') {
1056 key.pop();
1057 }
1058 if key.is_empty() {
1059 return Err("invalid cache key name".to_string());
1060 }
1061 Ok(key)
1062}
1063
1064fn secret_name_for_auth(auth_file: &Path, secret_dir: &Path) -> Option<String> {
1065 let auth_key = auth::identity_key_from_auth_file(auth_file)
1066 .ok()
1067 .flatten()?;
1068 let entries = fs::read_dir(secret_dir).ok()?;
1069 for entry in entries.flatten() {
1070 let path = entry.path();
1071 if path.extension().and_then(|s| s.to_str()) != Some("json") {
1072 continue;
1073 }
1074 let candidate_key = match auth::identity_key_from_auth_file(&path).ok().flatten() {
1075 Some(value) => value,
1076 None => continue,
1077 };
1078 if candidate_key == auth_key {
1079 return secret_file_basename(&path).ok();
1080 }
1081 }
1082 None
1083}
1084
1085fn current_secret_basename(secret_files: &[PathBuf]) -> Option<String> {
1086 let auth_file = paths::resolve_auth_file()?;
1087 if !auth_file.is_file() {
1088 return None;
1089 }
1090
1091 let auth_hash = gemini_fs::sha256_file(&auth_file).ok();
1092 if let Some(auth_hash) = auth_hash.as_deref() {
1093 for secret_file in secret_files {
1094 if let Ok(secret_hash) = gemini_fs::sha256_file(secret_file)
1095 && secret_hash == auth_hash
1096 && let Ok(name) = secret_file_basename(secret_file)
1097 {
1098 return Some(name);
1099 }
1100 }
1101 }
1102
1103 let auth_key = auth::identity_key_from_auth_file(&auth_file).ok().flatten();
1104 if let Some(auth_key) = auth_key.as_deref() {
1105 for secret_file in secret_files {
1106 if let Ok(Some(candidate_key)) = auth::identity_key_from_auth_file(secret_file)
1107 && candidate_key == auth_key
1108 && let Ok(name) = secret_file_basename(secret_file)
1109 {
1110 return Some(name);
1111 }
1112 }
1113 }
1114
1115 None
1116}
1117
1118fn starship_cache_dir() -> Option<PathBuf> {
1119 let root = cache_root()?;
1120 Some(root.join("gemini").join("starship-rate-limits"))
1121}
1122
1123fn cache_root() -> Option<PathBuf> {
1124 if let Ok(path) = std::env::var("ZSH_CACHE_DIR")
1125 && !path.is_empty()
1126 {
1127 return Some(PathBuf::from(path));
1128 }
1129 let zdotdir = paths::resolve_zdotdir()?;
1130 Some(zdotdir.join("cache"))
1131}
1132
1133fn render_line_for_summary(
1134 name: &str,
1135 summary: &RateLimitSummary,
1136 one_line: bool,
1137 time_format: &str,
1138) -> String {
1139 let reset = render::format_epoch_local(summary.weekly_reset_epoch, time_format)
1140 .unwrap_or_else(|| "?".to_string());
1141 let token_5h = format!(
1142 "{}:{}%",
1143 summary.non_weekly_label, summary.non_weekly_remaining
1144 );
1145 let token_weekly = format!("W:{}%", summary.weekly_remaining);
1146
1147 if one_line {
1148 return format!("{name} {token_5h} {token_weekly} {reset}");
1149 }
1150 format!("{token_5h} {token_weekly} {reset}")
1151}
1152
1153fn print_rate_limits_remaining(summary: &RateLimitSummary, time_format: &str) {
1154 println!("Rate limits remaining");
1155 let non_weekly_reset = summary
1156 .non_weekly_reset_epoch
1157 .and_then(|epoch| render::format_epoch_local(epoch, time_format))
1158 .unwrap_or_else(|| "?".to_string());
1159 let weekly_reset = render::format_epoch_local(summary.weekly_reset_epoch, time_format)
1160 .unwrap_or_else(|| "?".to_string());
1161 println!(
1162 "{} {}% • {}",
1163 summary.non_weekly_label, summary.non_weekly_remaining, non_weekly_reset
1164 );
1165 println!("Weekly {}% • {}", summary.weekly_remaining, weekly_reset);
1166}
1167
1168fn target_file_name(path: &Path) -> String {
1169 if let Some(name) = path.file_name().and_then(|value| value.to_str()) {
1170 name.to_string()
1171 } else {
1172 path.to_string_lossy().to_string()
1173 }
1174}
1175
1176fn now_epoch_seconds() -> i64 {
1177 std::time::SystemTime::now()
1178 .duration_since(std::time::UNIX_EPOCH)
1179 .map(|duration| duration.as_secs() as i64)
1180 .unwrap_or(0)
1181}
1182
1183fn env_non_empty(key: &str) -> Option<String> {
1184 std::env::var(key)
1185 .ok()
1186 .map(|raw| raw.trim().to_string())
1187 .filter(|raw| !raw.is_empty())
1188}
1189
1190fn env_truthy(key: &str) -> bool {
1191 match std::env::var(key) {
1192 Ok(raw) => {
1193 matches!(
1194 raw.trim().to_ascii_lowercase().as_str(),
1195 "1" | "true" | "yes" | "on"
1196 )
1197 }
1198 Err(_) => false,
1199 }
1200}
1201
1202struct RunError {
1203 code: String,
1204 message: String,
1205 details: Option<String>,
1206 exit_code: i32,
1207}
1208
1209fn emit_single_envelope(mode: &str, ok: bool, result: &JsonResultItem) {
1210 println!(
1211 "{{\"schema_version\":\"{}\",\"command\":\"{}\",\"mode\":\"{}\",\"ok\":{},\"result\":{}}}",
1212 DIAG_SCHEMA_VERSION,
1213 DIAG_COMMAND,
1214 json_escape(mode),
1215 if ok { "true" } else { "false" },
1216 result.to_json()
1217 );
1218}
1219
1220fn emit_collection_envelope(mode: &str, ok: bool, results: &[JsonResultItem]) {
1221 let mut body = String::new();
1222 body.push('[');
1223 for (index, result) in results.iter().enumerate() {
1224 if index > 0 {
1225 body.push(',');
1226 }
1227 body.push_str(&result.to_json());
1228 }
1229 body.push(']');
1230
1231 println!(
1232 "{{\"schema_version\":\"{}\",\"command\":\"{}\",\"mode\":\"{}\",\"ok\":{},\"results\":{}}}",
1233 DIAG_SCHEMA_VERSION,
1234 DIAG_COMMAND,
1235 json_escape(mode),
1236 if ok { "true" } else { "false" },
1237 body
1238 );
1239}
1240
1241fn emit_error_json(code: &str, message: &str, details: Option<String>) {
1242 print!(
1243 "{{\"schema_version\":\"{}\",\"command\":\"{}\",\"ok\":false,\"error\":{{\"code\":\"{}\",\"message\":\"{}\"",
1244 DIAG_SCHEMA_VERSION,
1245 DIAG_COMMAND,
1246 json_escape(code),
1247 json_escape(message),
1248 );
1249 if let Some(details) = details {
1250 print!(",\"details\":{}", details);
1251 }
1252 println!("}}}}");
1253}
1254
1255impl JsonResultItem {
1256 fn to_json(&self) -> String {
1257 let mut s = String::new();
1258 s.push('{');
1259 push_field(&mut s, "name", &json_string(&self.name), true);
1260 push_field(
1261 &mut s,
1262 "target_file",
1263 &json_string(&self.target_file),
1264 false,
1265 );
1266 push_field(&mut s, "status", &json_string(&self.status), false);
1267 push_field(&mut s, "ok", if self.ok { "true" } else { "false" }, false);
1268 push_field(&mut s, "source", &json_string(&self.source), false);
1269
1270 if let Some(summary) = &self.summary {
1271 push_field(&mut s, "summary", &summary.to_json(), false);
1272 }
1273
1274 if let Some(raw_usage) = &self.raw_usage {
1275 let trimmed = raw_usage.trim();
1276 if trimmed.starts_with('{') && trimmed.ends_with('}') {
1277 push_field(&mut s, "raw_usage", trimmed, false);
1278 } else {
1279 push_field(&mut s, "raw_usage", &json_string(trimmed), false);
1280 }
1281 } else {
1282 push_field(&mut s, "raw_usage", "null", false);
1283 }
1284
1285 if let (Some(code), Some(message)) = (&self.error_code, &self.error_message) {
1286 let error_json = format!(
1287 "{{\"code\":\"{}\",\"message\":\"{}\"}}",
1288 json_escape(code),
1289 json_escape(message)
1290 );
1291 push_field(&mut s, "error", &error_json, false);
1292 }
1293
1294 s.push('}');
1295 s
1296 }
1297}
1298
1299impl RateLimitSummary {
1300 fn to_json(&self) -> String {
1301 let mut s = String::new();
1302 s.push('{');
1303 push_field(
1304 &mut s,
1305 "non_weekly_label",
1306 &json_string(&self.non_weekly_label),
1307 true,
1308 );
1309 push_field(
1310 &mut s,
1311 "non_weekly_remaining",
1312 &self.non_weekly_remaining.to_string(),
1313 false,
1314 );
1315 match self.non_weekly_reset_epoch {
1316 Some(value) => push_field(&mut s, "non_weekly_reset_epoch", &value.to_string(), false),
1317 None => push_field(&mut s, "non_weekly_reset_epoch", "null", false),
1318 }
1319 push_field(
1320 &mut s,
1321 "weekly_remaining",
1322 &self.weekly_remaining.to_string(),
1323 false,
1324 );
1325 push_field(
1326 &mut s,
1327 "weekly_reset_epoch",
1328 &self.weekly_reset_epoch.to_string(),
1329 false,
1330 );
1331 s.push('}');
1332 s
1333 }
1334}
1335
1336fn push_field(buf: &mut String, key: &str, value_json: &str, first: bool) {
1337 if !first {
1338 buf.push(',');
1339 }
1340 buf.push('"');
1341 buf.push_str(&json_escape(key));
1342 buf.push_str("\":");
1343 buf.push_str(value_json);
1344}
1345
1346fn json_string(raw: &str) -> String {
1347 format!("\"{}\"", json_escape(raw))
1348}
1349
1350fn json_array(values: Vec<String>) -> String {
1351 let mut out = String::from("[");
1352 for (index, value) in values.iter().enumerate() {
1353 if index > 0 {
1354 out.push(',');
1355 }
1356 out.push_str(value);
1357 }
1358 out.push(']');
1359 out
1360}
1361
1362fn json_obj(fields: Vec<(String, String)>) -> String {
1363 let mut out = String::from("{");
1364 for (index, (key, value)) in fields.iter().enumerate() {
1365 if index > 0 {
1366 out.push(',');
1367 }
1368 out.push('"');
1369 out.push_str(&json_escape(key));
1370 out.push_str("\":");
1371 out.push_str(value);
1372 }
1373 out.push('}');
1374 out
1375}
1376
1377fn json_escape(raw: &str) -> String {
1378 let mut escaped = String::with_capacity(raw.len());
1379 for ch in raw.chars() {
1380 match ch {
1381 '"' => escaped.push_str("\\\""),
1382 '\\' => escaped.push_str("\\\\"),
1383 '\u{08}' => escaped.push_str("\\b"),
1384 '\u{0C}' => escaped.push_str("\\f"),
1385 '\n' => escaped.push_str("\\n"),
1386 '\r' => escaped.push_str("\\r"),
1387 '\t' => escaped.push_str("\\t"),
1388 ch if ch.is_control() => escaped.push_str(&format!("\\u{:04x}", ch as u32)),
1389 ch => escaped.push(ch),
1390 }
1391 }
1392 escaped
1393}
1394
1395#[cfg(test)]
1396mod tests {
1397 use super::*;
1398 use std::ffi::OsString;
1399
1400 struct EnvGuard {
1401 key: &'static str,
1402 old: Option<OsString>,
1403 }
1404
1405 impl EnvGuard {
1406 fn set(key: &'static str, value: impl AsRef<std::ffi::OsStr>) -> Self {
1407 let old = std::env::var_os(key);
1408 unsafe { std::env::set_var(key, value) };
1410 Self { key, old }
1411 }
1412 }
1413
1414 impl Drop for EnvGuard {
1415 fn drop(&mut self) {
1416 if let Some(value) = self.old.take() {
1417 unsafe { std::env::set_var(self.key, value) };
1419 } else {
1420 unsafe { std::env::remove_var(self.key) };
1422 }
1423 }
1424 }
1425
1426 #[test]
1427 fn cache_key_normalizes_and_rejects_empty() {
1428 assert_eq!(cache_key("Alpha.Work").expect("key"), "alpha_work");
1429 assert!(cache_key("___").is_err());
1430 }
1431
1432 #[test]
1433 fn secret_file_basename_requires_non_empty_name() {
1434 assert_eq!(
1435 secret_file_basename(Path::new("/tmp/alpha.json")).expect("basename"),
1436 "alpha"
1437 );
1438 assert!(secret_file_basename(Path::new("/tmp/.json")).is_err());
1439 }
1440
1441 #[test]
1442 fn env_truthy_accepts_expected_variants() {
1443 let _v1 = EnvGuard::set("GEMINI_TEST_TRUTHY", "true");
1444 assert!(env_truthy("GEMINI_TEST_TRUTHY"));
1445 let _v2 = EnvGuard::set("GEMINI_TEST_TRUTHY", "ON");
1446 assert!(env_truthy("GEMINI_TEST_TRUTHY"));
1447 let _v3 = EnvGuard::set("GEMINI_TEST_TRUTHY", "0");
1448 assert!(!env_truthy("GEMINI_TEST_TRUTHY"));
1449 }
1450
1451 #[test]
1452 fn render_line_for_summary_formats_name_and_one_line() {
1453 let summary = RateLimitSummary {
1454 non_weekly_label: "5h".to_string(),
1455 non_weekly_remaining: 94,
1456 non_weekly_reset_epoch: Some(1700003600),
1457 weekly_remaining: 88,
1458 weekly_reset_epoch: 1700600000,
1459 };
1460 assert_eq!(
1461 render_line_for_summary("alpha", &summary, false, "%m-%d %H:%M"),
1462 "5h:94% W:88% 11-21 20:53"
1463 );
1464 assert_eq!(
1465 render_line_for_summary("alpha", &summary, true, "%m-%d %H:%M"),
1466 "alpha 5h:94% W:88% 11-21 20:53"
1467 );
1468 }
1469
1470 #[test]
1471 fn json_helpers_escape_and_build_structures() {
1472 assert_eq!(json_escape("a\"b\\n"), "a\\\"b\\\\n");
1473 assert_eq!(
1474 json_array(vec![json_string("a"), json_string("b")]),
1475 "[\"a\",\"b\"]"
1476 );
1477 assert_eq!(
1478 json_obj(vec![
1479 ("k1".to_string(), json_string("v1")),
1480 ("k2".to_string(), "2".to_string())
1481 ]),
1482 "{\"k1\":\"v1\",\"k2\":2}"
1483 );
1484 }
1485
1486 #[test]
1487 fn rate_limit_summary_to_json_includes_null_non_weekly_reset() {
1488 let summary = RateLimitSummary {
1489 non_weekly_label: "5h".to_string(),
1490 non_weekly_remaining: 90,
1491 non_weekly_reset_epoch: None,
1492 weekly_remaining: 80,
1493 weekly_reset_epoch: 1700600000,
1494 };
1495 let rendered = summary.to_json();
1496 assert!(rendered.contains("\"non_weekly_reset_epoch\":null"));
1497 assert!(rendered.contains("\"weekly_reset_epoch\":1700600000"));
1498 }
1499
1500 #[test]
1501 fn json_result_item_to_json_supports_error_and_raw_usage_variants() {
1502 let item = JsonResultItem {
1503 name: "alpha".to_string(),
1504 target_file: "alpha.json".to_string(),
1505 status: "error".to_string(),
1506 ok: false,
1507 source: "network".to_string(),
1508 summary: None,
1509 raw_usage: Some("{\"rate_limit\":{}}".to_string()),
1510 error_code: Some("request-failed".to_string()),
1511 error_message: Some("boom".to_string()),
1512 };
1513 let rendered = item.to_json();
1514 assert!(rendered.contains("\"raw_usage\":{\"rate_limit\":{}}"));
1515 assert!(rendered.contains("\"error\":{\"code\":\"request-failed\",\"message\":\"boom\"}"));
1516 }
1517
1518 #[test]
1519 fn collect_secret_files_returns_sorted_json_files() {
1520 let dir = tempfile::TempDir::new().expect("tempdir");
1521 let secrets = dir.path().join("secrets");
1522 std::fs::create_dir_all(&secrets).expect("secrets");
1523 std::fs::write(secrets.join("b.json"), "{}").expect("b");
1524 std::fs::write(secrets.join("a.json"), "{}").expect("a");
1525 std::fs::write(secrets.join("skip.txt"), "x").expect("skip");
1526 let _secret = EnvGuard::set("GEMINI_SECRET_DIR", &secrets);
1527
1528 let files = collect_secret_files().expect("files");
1529 assert_eq!(
1530 files
1531 .iter()
1532 .map(|p| p.file_name().and_then(|v| v.to_str()).unwrap_or_default())
1533 .collect::<Vec<_>>(),
1534 vec!["a.json", "b.json"]
1535 );
1536 }
1537
1538 #[test]
1539 fn resolve_single_target_appends_json_when_secret_dir_is_configured() {
1540 let dir = tempfile::TempDir::new().expect("tempdir");
1541 let secrets = dir.path().join("secrets");
1542 std::fs::create_dir_all(&secrets).expect("secrets");
1543 let target = secrets.join("alpha.json");
1544 std::fs::write(&target, "{}").expect("target");
1545 let _secret = EnvGuard::set("GEMINI_SECRET_DIR", &secrets);
1546
1547 let resolved = resolve_single_target(Some("alpha")).expect("resolved");
1548 assert_eq!(resolved, target);
1549 }
1550
1551 #[test]
1552 fn clear_starship_cache_rejects_non_absolute_cache_root() {
1553 let _cache = EnvGuard::set("ZSH_CACHE_DIR", "relative-cache");
1554 let err = clear_starship_cache().expect_err("non-absolute should fail");
1555 assert!(err.contains("non-absolute cache root"));
1556 }
1557
1558 #[test]
1559 fn emit_helpers_cover_single_collection_and_error_envelopes() {
1560 let item = JsonResultItem {
1561 name: "alpha".to_string(),
1562 target_file: "alpha.json".to_string(),
1563 status: "ok".to_string(),
1564 ok: true,
1565 source: "network".to_string(),
1566 summary: Some(RateLimitSummary {
1567 non_weekly_label: "5h".to_string(),
1568 non_weekly_remaining: 94,
1569 non_weekly_reset_epoch: Some(1700003600),
1570 weekly_remaining: 88,
1571 weekly_reset_epoch: 1700600000,
1572 }),
1573 raw_usage: Some("{\"rate_limit\":{}}".to_string()),
1574 error_code: None,
1575 error_message: None,
1576 };
1577 emit_single_envelope("single", true, &item);
1578 emit_collection_envelope("all", true, &[item]);
1579 emit_error_json("failure", "boom", None);
1580 }
1581}