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