1use std::collections::BTreeSet;
7use std::sync::{Mutex, OnceLock};
8use std::time::Duration;
9
10use clap::ValueEnum;
11
12use crate::usage::VendorSnapshot;
13use crate::widget::cli::Cli;
14
15pub const HTTP_CLIENT_TIMEOUT: Duration = Duration::from_secs(30);
18
19pub const MAX_BODY_BYTES: usize = 2 * 1024 * 1024;
24
25pub(crate) const VENDOR_SECRET_ENV_VARS: &[&str] = &[
28 "ZAI_API_KEY",
29 "OPENROUTER_API_KEY",
30 "DEEPSEEK_API_KEY",
31 "KIMI_API_KEY",
32 "KILO_API_KEY",
33 "NOVITA_API_KEY",
34 "MINIMAX_API_KEY",
35 "MOONSHOT_API_KEY",
36 "XAI_MANAGEMENT_KEY",
37 "ANTHROPIC_ADMIN_KEY",
38 "XAI_API_KEY",
39 "GROK_API_KEY",
40 "OPENCODE_GO_API_KEY",
41 "COMMANDCODE_API_KEY",
42 "ORCAROUTER_API_KEY",
43 "GITHUB_COPILOT_TOKEN",
44 "GH_TOKEN",
45 "GITHUB_TOKEN",
46 "OLLAMA_API_KEY",
47];
48
49fn registered_secret_env_vars() -> &'static Mutex<BTreeSet<&'static str>> {
54 static REGISTERED: OnceLock<Mutex<BTreeSet<&'static str>>> = OnceLock::new();
55 REGISTERED.get_or_init(|| Mutex::new(BTreeSet::new()))
56}
57
58pub fn register_secret_env_vars(names: &[String]) {
69 let mut registered = registered_secret_env_vars()
70 .lock()
71 .unwrap_or_else(|poisoned| poisoned.into_inner());
72 for name in names {
73 if !crate::config::is_valid_env_var_name(name)
74 || VENDOR_SECRET_ENV_VARS.contains(&name.as_str())
75 || registered.contains(name.as_str())
76 {
77 continue;
78 }
79 registered.insert(Box::leak(name.clone().into_boxed_str()));
80 }
81}
82
83pub(crate) fn vendor_secret_env_vars_to_remove(keep: &[&str]) -> Vec<&'static str> {
84 let registered = registered_secret_env_vars()
85 .lock()
86 .unwrap_or_else(|poisoned| poisoned.into_inner());
87 VENDOR_SECRET_ENV_VARS
88 .iter()
89 .copied()
90 .chain(registered.iter().copied())
91 .filter(|var| !keep.contains(var))
92 .collect()
93}
94
95pub fn same_origin_redirect_policy() -> reqwest::redirect::Policy {
100 reqwest::redirect::Policy::custom(|attempt| {
101 if attempt.previous().len() >= 10 {
102 return attempt.error("too many redirects");
103 }
104 let Some(origin) = attempt.previous().first() else {
105 return attempt.stop();
106 };
107 let target = attempt.url();
108 if target.scheme() == origin.scheme()
109 && target.host_str() == origin.host_str()
110 && target.port_or_known_default() == origin.port_or_known_default()
111 {
112 attempt.follow()
113 } else {
114 attempt.stop()
115 }
116 })
117}
118
119pub async fn read_body_capped(
127 mut resp: reqwest::Response,
128 max: usize,
129) -> crate::error::Result<Vec<u8>> {
130 let too_big = |n: u64| {
131 crate::error::AppError::Schema(format!(
132 "response body exceeds the {max}-byte limit ({n} bytes); refusing to buffer it"
133 ))
134 };
135 if let Some(len) = resp.content_length()
136 && len > max as u64
137 {
138 return Err(too_big(len));
139 }
140 let mut buf: Vec<u8> = Vec::new();
141 while let Some(chunk) = resp.chunk().await? {
142 if chunk.len() > max.saturating_sub(buf.len()) {
143 return Err(too_big(buf.len().saturating_add(chunk.len()) as u64));
144 }
145 buf.extend_from_slice(&chunk);
146 }
147 Ok(buf)
148}
149
150#[derive(
152 Debug, Clone, Copy, ValueEnum, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize,
153)]
154#[serde(rename_all = "lowercase")]
155pub enum VendorId {
156 Anthropic,
157 #[serde(rename = "anthropic_api")]
158 AnthropicApi,
159 Openai,
160 Copilot,
161 Zai,
162 Openrouter,
163 Deepseek,
164 Kimi,
165 Kilo,
166 Novita,
167 Moonshot,
168 Grok,
169 Supergrok,
170 Grokbot,
171 Antigravity,
172 Cursor,
173 Minimax,
174 Kiro,
175 #[serde(rename = "nous")]
176 NousResearch,
177 #[serde(rename = "opencode-go")]
178 OpenCodeGo,
179 #[serde(rename = "commandcode")]
180 CommandCode,
181 Ollama,
182 OrcaRouter,
183 ModelStudio,
184}
185
186#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
190#[serde(rename_all = "lowercase")]
191pub enum AuthKind {
192 Oauth,
194 ApiKey,
196 Local,
199}
200
201impl AuthKind {
202 pub const fn as_str(self) -> &'static str {
203 match self {
204 AuthKind::Oauth => "oauth",
205 AuthKind::ApiKey => "apikey",
206 AuthKind::Local => "local",
207 }
208 }
209}
210
211impl VendorId {
212 pub fn slug(self) -> &'static str {
213 match self {
214 VendorId::Anthropic => "anthropic",
215 VendorId::AnthropicApi => "anthropic_api",
216 VendorId::Openai => "openai",
217 VendorId::Copilot => "copilot",
218 VendorId::Zai => "zai",
219 VendorId::Openrouter => "openrouter",
220 VendorId::Deepseek => "deepseek",
221 VendorId::Kimi => "kimi",
222 VendorId::Kilo => "kilo",
223 VendorId::Novita => "novita",
224 VendorId::Moonshot => "moonshot",
225 VendorId::Grok => "grok",
226 VendorId::Supergrok => "supergrok",
227 VendorId::Grokbot => "grokbot",
228 VendorId::Antigravity => "antigravity",
229 VendorId::Cursor => "cursor",
230 VendorId::Minimax => "minimax",
231 VendorId::Kiro => "kiro",
232 VendorId::NousResearch => "nous",
233 VendorId::OpenCodeGo => "opencode-go",
234 VendorId::CommandCode => "commandcode",
235 VendorId::Ollama => "ollama",
236 VendorId::OrcaRouter => "orcarouter",
237 VendorId::ModelStudio => "modelstudio",
238 }
239 }
240
241 pub fn display_name(self) -> &'static str {
245 match self {
246 VendorId::Anthropic => "Claude",
247 VendorId::AnthropicApi => "Anthropic API",
248 VendorId::Openai => "Codex",
249 VendorId::Copilot => "GitHub Copilot",
250 VendorId::Zai => "Z.AI",
251 VendorId::Openrouter => "OpenRouter",
252 VendorId::Deepseek => "DeepSeek",
253 VendorId::Kimi => "Kimi",
254 VendorId::Kilo => "Kilo",
255 VendorId::Novita => "Novita",
256 VendorId::Moonshot => "Moonshot",
257 VendorId::Grok => "Grok",
258 VendorId::Supergrok => "SuperGrok",
259 VendorId::Grokbot => "Grok Bot",
260 VendorId::Antigravity => "Antigravity",
261 VendorId::Cursor => "Cursor",
262 VendorId::Minimax => "MiniMax",
263 VendorId::Kiro => "Kiro",
264 VendorId::NousResearch => "Nous Research",
265 VendorId::OpenCodeGo => "OpenCode Go",
266 VendorId::CommandCode => "Command Code",
267 VendorId::Ollama => "Ollama Cloud",
268 VendorId::OrcaRouter => "OrcaRouter",
269 VendorId::ModelStudio => "Model Studio",
270 }
271 }
272
273 pub const fn bar_icon(self) -> &'static str {
277 match self {
278 VendorId::Anthropic => "",
279 VendorId::AnthropicApi => "",
280 VendorId::Openai => "",
281 VendorId::Copilot => "",
282 VendorId::Zai => VendorId::Zai.short_name(),
283 VendorId::Openrouter => "",
284 VendorId::Deepseek => "",
285 VendorId::Kimi => VendorId::Kimi.short_name(),
286 VendorId::Kilo => "",
287 VendorId::Novita => "",
288 VendorId::Moonshot => VendorId::Moonshot.short_name(),
289 VendorId::Grok | VendorId::Supergrok => "",
290 VendorId::Grokbot => VendorId::Grokbot.short_name(),
291 VendorId::Antigravity => VendorId::Antigravity.short_name(),
292 VendorId::Cursor => "❯",
293 VendorId::Minimax => VendorId::Minimax.short_name(),
294 VendorId::Kiro => "◆",
295 VendorId::NousResearch => VendorId::NousResearch.short_name(),
296 VendorId::OpenCodeGo => VendorId::OpenCodeGo.short_name(),
297 VendorId::CommandCode => VendorId::CommandCode.short_name(),
298 VendorId::Ollama => VendorId::Ollama.short_name(),
301 VendorId::OrcaRouter => VendorId::OrcaRouter.short_name(),
303 VendorId::ModelStudio => VendorId::ModelStudio.short_name(),
305 }
306 }
307
308 pub const fn short_name(self) -> &'static str {
313 match self {
314 VendorId::Anthropic => "cld",
315 VendorId::AnthropicApi => "aac",
316 VendorId::Openai => "gpt",
317 VendorId::Copilot => "ghc",
318 VendorId::Zai => "zai",
319 VendorId::Openrouter => "opr",
320 VendorId::Deepseek => "dsk",
321 VendorId::Kimi => "kmi",
322 VendorId::Kilo => "klo",
323 VendorId::Novita => "nvt",
324 VendorId::Moonshot => "msh",
325 VendorId::Grok => "grk",
326 VendorId::Supergrok => "sgk",
327 VendorId::Grokbot => "gbt",
328 VendorId::Antigravity => "agy",
329 VendorId::Cursor => "cur",
330 VendorId::Minimax => "mmx",
331 VendorId::Kiro => "kir",
332 VendorId::NousResearch => "nrs",
333 VendorId::OpenCodeGo => "ocg",
334 VendorId::CommandCode => "cmc",
335 VendorId::Ollama => "oll",
336 VendorId::OrcaRouter => "orc",
337 VendorId::ModelStudio => "mst",
338 }
339 }
340
341 pub const fn config_section(self) -> &'static str {
349 match self {
350 VendorId::Anthropic => "anthropic",
351 VendorId::AnthropicApi => "anthropic_api",
352 VendorId::Openai => "openai",
353 VendorId::Copilot => "copilot",
354 VendorId::Zai => "zai",
355 VendorId::Openrouter => "openrouter",
356 VendorId::Deepseek => "deepseek",
357 VendorId::Kimi => "kimi",
358 VendorId::Kilo => "kilo",
359 VendorId::Novita => "novita",
360 VendorId::Moonshot => "moonshot",
361 VendorId::Grok => "grok",
362 VendorId::Supergrok => "supergrok",
363 VendorId::Grokbot => "grokbot",
364 VendorId::Antigravity => "antigravity",
365 VendorId::Cursor => "cursor",
366 VendorId::Minimax => "minimax",
367 VendorId::Kiro => "kiro",
368 VendorId::NousResearch => "nous",
369 VendorId::OpenCodeGo => "opencode-go",
370 VendorId::CommandCode => "commandcode",
371 VendorId::Ollama => "ollama",
372 VendorId::OrcaRouter => "orcarouter",
373 VendorId::ModelStudio => "modelstudio",
374 }
375 }
376
377 pub const fn auth_kind(self) -> AuthKind {
383 match self {
384 VendorId::Anthropic
385 | VendorId::Openai
386 | VendorId::Copilot
387 | VendorId::NousResearch
388 | VendorId::CommandCode => AuthKind::Oauth,
389 VendorId::AnthropicApi
390 | VendorId::Zai
391 | VendorId::Openrouter
392 | VendorId::Deepseek
393 | VendorId::Kimi
394 | VendorId::Kilo
395 | VendorId::Novita
396 | VendorId::Moonshot
397 | VendorId::Grok
398 | VendorId::Minimax
399 | VendorId::OpenCodeGo
400 | VendorId::Ollama
401 | VendorId::OrcaRouter => AuthKind::ApiKey,
402 VendorId::Supergrok
410 | VendorId::Antigravity
411 | VendorId::Cursor
412 | VendorId::Kiro
413 | VendorId::Grokbot
414 | VendorId::ModelStudio => AuthKind::Local,
415 }
416 }
417
418 pub const fn api_key_env(self) -> &'static str {
423 match self {
424 VendorId::AnthropicApi => "ANTHROPIC_ADMIN_KEY",
425 VendorId::Zai => "ZAI_API_KEY",
426 VendorId::Openrouter => "OPENROUTER_API_KEY",
427 VendorId::Deepseek => "DEEPSEEK_API_KEY",
428 VendorId::Kimi => "KIMI_API_KEY",
429 VendorId::Kilo => "KILO_API_KEY",
430 VendorId::Novita => "NOVITA_API_KEY",
431 VendorId::Moonshot => "MOONSHOT_API_KEY",
432 VendorId::Grok => "XAI_MANAGEMENT_KEY",
433 VendorId::Minimax => "MINIMAX_API_KEY",
434 VendorId::OpenCodeGo => "OPENCODE_GO_API_KEY",
435 VendorId::Ollama => "OLLAMA_API_KEY",
436 VendorId::OrcaRouter => "ORCAROUTER_API_KEY",
437 VendorId::Copilot => "GITHUB_COPILOT_TOKEN",
441 VendorId::CommandCode => "COMMANDCODE_API_KEY",
442 VendorId::Anthropic
443 | VendorId::Openai
444 | VendorId::Supergrok
445 | VendorId::Grokbot
446 | VendorId::Antigravity
447 | VendorId::Cursor
448 | VendorId::Kiro
449 | VendorId::NousResearch
450 | VendorId::ModelStudio => "",
451 }
452 }
453
454 pub const fn sign_in_hint(self) -> &'static str {
467 match self {
468 VendorId::Anthropic => "Run `claude` in a terminal, then Refresh.",
469 VendorId::Openai => "Run `codex login` in a terminal, then Refresh.",
470 VendorId::Copilot => "Run `gh auth login` in a terminal, then Refresh.",
471 VendorId::Kiro => "Run `kiro-cli login` in a terminal, then Refresh.",
472 VendorId::Kimi => "Run `kimi` in a terminal, or set an API key.",
473 VendorId::CommandCode => "Run `commandcode` in a terminal, then Refresh.",
474 VendorId::NousResearch => {
475 "Run `ai-usagebar auth nous login` in a terminal, then Refresh."
476 }
477 VendorId::Cursor => "Sign in to the Cursor app, then Refresh.",
478 VendorId::Antigravity => "Open Antigravity or run `agy`, then Refresh.",
479 VendorId::Grok | VendorId::Supergrok => "Sign in with `grok`, then Refresh.",
480 VendorId::Grokbot => "Install and sign in to the Grok Bot desktop app, then Refresh.",
481 VendorId::ModelStudio => {
483 "Install the official `bl` CLI and run `bl auth login --console`, then Refresh."
484 }
485 VendorId::AnthropicApi
490 | VendorId::Zai
491 | VendorId::Openrouter
492 | VendorId::Deepseek
493 | VendorId::Kilo
494 | VendorId::Novita
495 | VendorId::Moonshot
496 | VendorId::Minimax
497 | VendorId::OpenCodeGo
498 | VendorId::Ollama
499 | VendorId::OrcaRouter => "Add an API key in Settings, then Refresh.",
500 }
501 }
502
503 pub const fn login_command(self) -> &'static str {
504 match self {
505 VendorId::Anthropic => "claude",
506 VendorId::Openai => "codex login",
507 VendorId::Copilot => "gh auth login",
508 VendorId::CommandCode => "commandcode",
509 VendorId::NousResearch => "ai-usagebar auth nous login",
510 VendorId::Kiro => "kiro-cli login",
511 VendorId::ModelStudio => "bl auth login --console",
513 VendorId::Kimi => "kimi",
516 VendorId::AnthropicApi
517 | VendorId::Zai
518 | VendorId::Openrouter
519 | VendorId::Deepseek
520 | VendorId::Kilo
521 | VendorId::Novita
522 | VendorId::Moonshot
523 | VendorId::Grok
524 | VendorId::Supergrok
525 | VendorId::Grokbot
526 | VendorId::Antigravity
527 | VendorId::Cursor
528 | VendorId::Minimax
529 | VendorId::OpenCodeGo
530 | VendorId::Ollama
531 | VendorId::OrcaRouter => "",
532 }
533 }
534
535 pub const fn all() -> &'static [VendorId] {
536 &[
537 VendorId::Anthropic,
538 VendorId::AnthropicApi,
539 VendorId::Openai,
540 VendorId::Copilot,
541 VendorId::Zai,
542 VendorId::Openrouter,
543 VendorId::Deepseek,
544 VendorId::Kimi,
545 VendorId::Kilo,
546 VendorId::Novita,
547 VendorId::Moonshot,
548 VendorId::Grok,
549 VendorId::Supergrok,
550 VendorId::Grokbot,
551 VendorId::Antigravity,
552 VendorId::Cursor,
553 VendorId::Minimax,
554 VendorId::Kiro,
555 VendorId::NousResearch,
556 VendorId::OpenCodeGo,
557 VendorId::CommandCode,
558 VendorId::Ollama,
559 VendorId::OrcaRouter,
560 VendorId::ModelStudio,
561 ]
562 }
563
564 pub fn from_slug(slug: &str) -> Option<Self> {
569 Self::all().iter().copied().find(|id| id.slug() == slug)
570 }
571}
572
573pub type VendorOutcome = crate::outcome::Outcome<VendorSnapshot>;
578
579#[derive(Debug, Clone)]
581pub struct RenderOpts {
582 pub format: Option<String>,
583 pub tooltip_format: Option<String>,
584 pub icon: Option<String>,
585 pub pace_tolerance: u32,
586 pub format_pace_color: bool,
587 pub tooltip_pace_pts: bool,
588}
589
590impl RenderOpts {
591 pub fn from_cli(cli: &Cli) -> Self {
592 Self {
593 format: cli.format.clone(),
594 tooltip_format: cli.tooltip_format.clone(),
595 icon: cli.icon.clone(),
596 pace_tolerance: cli.pace_tolerance,
597 format_pace_color: cli.format_pace_color,
598 tooltip_pace_pts: cli.tooltip_pace_pts,
599 }
600 }
601}
602
603#[cfg(test)]
604mod tests {
605 use super::*;
606
607 #[test]
608 fn every_vendor_has_stable_machine_and_display_names() {
609 for vendor in VendorId::all() {
610 assert!(!vendor.slug().is_empty());
611 assert!(!vendor.display_name().is_empty());
612 }
613 assert_eq!(VendorId::Anthropic.slug(), "anthropic");
614 assert_eq!(VendorId::Anthropic.display_name(), "Claude");
615 assert_eq!(VendorId::Openai.display_name(), "Codex");
616 assert_eq!(VendorId::Zai.display_name(), "Z.AI");
617 }
618
619 #[test]
623 fn every_vendor_short_name_is_a_unique_three_letter_code() {
624 let mut seen = std::collections::BTreeSet::new();
625 for vendor in VendorId::all() {
626 let short = vendor.short_name();
627 assert_eq!(short.len(), 3, "{} is not three letters", vendor.slug());
628 assert!(
629 short.chars().all(|c| c.is_ascii_lowercase()),
630 "{} is not lowercase ascii",
631 vendor.slug()
632 );
633 assert!(seen.insert(short), "{short} is used by two vendors");
634 }
635 assert_eq!(VendorId::Anthropic.short_name(), "cld");
636 assert_eq!(VendorId::Openai.short_name(), "gpt");
637 assert_eq!(VendorId::Zai.short_name(), "zai");
638 assert_eq!(VendorId::Antigravity.short_name(), "agy");
639 }
640
641 #[test]
647 fn every_vendor_has_a_bar_icon_and_no_two_share_one() {
648 use std::collections::BTreeMap;
649
650 let mut by_icon: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
651 for vendor in VendorId::all() {
652 assert!(!vendor.bar_icon().is_empty(), "{}", vendor.slug());
653 by_icon
654 .entry(vendor.bar_icon())
655 .or_default()
656 .push(vendor.slug());
657 }
658
659 let shared: Vec<_> = by_icon
660 .iter()
661 .filter(|(_, vendors)| vendors.len() > 1)
662 .filter(|(_, vendors)| vendors.as_slice() != ["grok", "supergrok"])
663 .collect();
664 assert!(
665 shared.is_empty(),
666 "these providers are indistinguishable in a bar that shows them \
667 side by side: {shared:#?}"
668 );
669 assert_eq!(VendorId::Anthropic.bar_icon(), "");
670 assert_eq!(VendorId::Openai.bar_icon(), "");
671 assert_eq!(VendorId::Supergrok.bar_icon(), VendorId::Grok.bar_icon());
672 assert_eq!(VendorId::CommandCode.bar_icon(), "cmc");
673 }
674
675 #[test]
676 fn new_vendor_contracts_keep_public_names_and_slugs() {
677 assert_eq!(VendorId::NousResearch.slug(), "nous");
678 assert_eq!(VendorId::NousResearch.display_name(), "Nous Research");
679 assert_eq!(VendorId::OpenCodeGo.slug(), "opencode-go");
680 assert_eq!(VendorId::OpenCodeGo.display_name(), "OpenCode Go");
681 assert_eq!(
682 serde_json::to_value(VendorId::OpenCodeGo).unwrap(),
683 serde_json::json!("opencode-go")
684 );
685 }
686
687 #[test]
688 fn vendor_secret_env_vars_cover_config_defaults() {
689 let configured_defaults = [
690 "ZAI_API_KEY",
691 "OPENROUTER_API_KEY",
692 "DEEPSEEK_API_KEY",
693 "KIMI_API_KEY",
694 "KILO_API_KEY",
695 "NOVITA_API_KEY",
696 "MINIMAX_API_KEY",
697 "MOONSHOT_API_KEY",
698 "XAI_MANAGEMENT_KEY",
699 "ANTHROPIC_ADMIN_KEY",
700 "GITHUB_COPILOT_TOKEN",
701 "ORCAROUTER_API_KEY",
702 ];
703 for name in configured_defaults {
704 assert!(VENDOR_SECRET_ENV_VARS.contains(&name), "missing {name}");
705 }
706 }
707
708 #[test]
709 fn vars_to_remove_preserves_only_requested_grok_credentials() {
710 let removed = vendor_secret_env_vars_to_remove(&["XAI_API_KEY", "GROK_API_KEY"]);
711 assert!(!removed.contains(&"XAI_API_KEY"));
712 assert!(!removed.contains(&"GROK_API_KEY"));
713 assert!(removed.contains(&"ANTHROPIC_ADMIN_KEY"));
714 assert!(removed.contains(&"OPENROUTER_API_KEY"));
715 let builtins = removed
718 .iter()
719 .filter(|var| VENDOR_SECRET_ENV_VARS.contains(var))
720 .count();
721 assert_eq!(builtins, VENDOR_SECRET_ENV_VARS.len() - 2);
722 }
723
724 #[test]
725 fn a_registered_custom_env_var_is_scrubbed_like_a_builtin_one() {
726 let name = "AI_USAGEBAR_TEST_CUSTOM_TOKEN_7F3A";
727 assert!(!vendor_secret_env_vars_to_remove(&[]).contains(&name));
728
729 register_secret_env_vars(&[name.to_string(), "not a name!".to_string()]);
730 register_secret_env_vars(&[name.to_string()]);
731
732 let removed = vendor_secret_env_vars_to_remove(&[]);
733 assert_eq!(
734 removed.iter().filter(|var| **var == name).count(),
735 1,
736 "registering twice must not list it twice: {removed:?}"
737 );
738 assert!(!removed.contains(&"not a name!"), "{removed:?}");
739 assert!(
740 !vendor_secret_env_vars_to_remove(&[name]).contains(&name),
741 "`keep` applies to registered names too"
742 );
743 }
744
745 #[test]
746 fn copilot_token_is_removed_before_unrelated_subprocesses_launch() {
747 let removed = vendor_secret_env_vars_to_remove(&[]);
748 assert!(removed.contains(&"GITHUB_COPILOT_TOKEN"));
749 }
750
751 #[tokio::test]
752 async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
753 let mut server = mockito::Server::new_async().await;
754 server
755 .mock("GET", "/big")
756 .with_status(200)
757 .with_body("x".repeat(4096))
758 .create_async()
759 .await;
760 server
761 .mock("GET", "/small")
762 .with_status(200)
763 .with_body("hello")
764 .create_async()
765 .await;
766
767 let client = reqwest::Client::new();
768
769 let resp = client
771 .get(format!("{}/big", server.url()))
772 .send()
773 .await
774 .unwrap();
775 let err = read_body_capped(resp, 1024).await.unwrap_err();
776 assert!(
777 err.to_string().contains("exceeds"),
778 "unexpected error: {err}"
779 );
780
781 let resp = client
783 .get(format!("{}/small", server.url()))
784 .send()
785 .await
786 .unwrap();
787 assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
788 }
789
790 #[tokio::test]
791 async fn chunked_body_without_content_length_still_hits_the_cap() {
792 let mut server = mockito::Server::new_async().await;
793 server
794 .mock("GET", "/chunked")
795 .with_status(200)
796 .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
797 .create_async()
798 .await;
799
800 let response = reqwest::Client::new()
801 .get(format!("{}/chunked", server.url()))
802 .send()
803 .await
804 .unwrap();
805 assert!(response.content_length().is_none());
806 let error = read_body_capped(response, 1024).await.unwrap_err();
807 assert!(error.to_string().contains("exceeds"), "{error}");
808 }
809
810 #[tokio::test]
811 async fn same_origin_redirects_still_work_with_vendor_headers() {
812 let mut server = mockito::Server::new_async().await;
813 let redirect = server
814 .mock("GET", "/start")
815 .match_header("x-api-key", "secret")
816 .with_status(302)
817 .with_header("location", "/finish")
818 .create_async()
819 .await;
820 let finish = server
821 .mock("GET", "/finish")
822 .match_header("x-api-key", "secret")
823 .with_status(200)
824 .create_async()
825 .await;
826 let client = reqwest::Client::builder()
827 .redirect(same_origin_redirect_policy())
828 .build()
829 .unwrap();
830
831 let response = client
832 .get(format!("{}/start", server.url()))
833 .header("x-api-key", "secret")
834 .send()
835 .await
836 .unwrap();
837
838 assert_eq!(response.status(), reqwest::StatusCode::OK);
839 redirect.assert_async().await;
840 finish.assert_async().await;
841 }
842
843 #[tokio::test]
844 async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
845 let mut origin = mockito::Server::new_async().await;
846 let mut target = mockito::Server::new_async().await;
847 let target_url = format!("{}/capture", target.url());
848 let redirect = origin
849 .mock("GET", "/start")
850 .match_header("x-api-key", "secret")
851 .with_status(302)
852 .with_header("location", &target_url)
853 .create_async()
854 .await;
855 let capture = target
856 .mock("GET", "/capture")
857 .expect(0)
858 .create_async()
859 .await;
860 let client = reqwest::Client::builder()
861 .redirect(same_origin_redirect_policy())
862 .build()
863 .unwrap();
864
865 let response = client
866 .get(format!("{}/start", origin.url()))
867 .header("x-api-key", "secret")
868 .send()
869 .await
870 .unwrap();
871
872 assert_eq!(response.status(), reqwest::StatusCode::FOUND);
873 redirect.assert_async().await;
874 capture.assert_async().await;
875 }
876
877 #[test]
878 fn vendor_id_slug_round_trip() {
879 for id in VendorId::all() {
880 assert_eq!(
881 id.slug(),
882 serde_json::to_value(id).unwrap().as_str().unwrap()
883 );
884 }
885 }
886
887 #[test]
891 fn from_slug_resolves_every_slug_and_rejects_everything_else() {
892 for id in VendorId::all() {
893 assert_eq!(VendorId::from_slug(id.slug()), Some(*id));
894 }
895 assert_eq!(
896 VendorId::from_slug("anthropic_api"),
897 Some(VendorId::AnthropicApi)
898 );
899 assert_eq!(
900 VendorId::from_slug("opencode-go"),
901 Some(VendorId::OpenCodeGo)
902 );
903 assert_eq!(VendorId::from_slug(""), None);
904 assert_eq!(VendorId::from_slug("custom"), None);
905 assert_eq!(VendorId::from_slug("Anthropic"), None);
906 }
907}