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