ai_usagebar/
vendor.rs

1//! Shared vendor IDs and renderer/fetcher structs used by the widget and TUI.
2//!
3//! Snapshots remain a discriminated `VendorSnapshot` enum because the vendors
4//! have genuinely different shapes — see `usage.rs`.
5
6use std::time::Duration;
7
8use clap::ValueEnum;
9
10use crate::usage::VendorSnapshot;
11use crate::widget::cli::Cli;
12
13/// Outer reqwest client timeout shared by widget and TUI entry points.
14/// Vendor fetchers still apply their own tighter per-request timeouts.
15pub const HTTP_CLIENT_TIMEOUT: Duration = Duration::from_secs(30);
16
17/// Upper bound on a vendor response body. Every one of these endpoints returns
18/// a small JSON document — the largest observed is a few kilobytes — so this is
19/// generous by three orders of magnitude while still bounding the damage from a
20/// misbehaving proxy or a hijacked endpoint.
21pub const MAX_BODY_BYTES: usize = 2 * 1024 * 1024;
22
23/// Credential-bearing environment variables owned by ai-usagebar vendors.
24/// Subprocesses receive only the entries that belong to their own provider.
25pub(crate) const VENDOR_SECRET_ENV_VARS: &[&str] = &[
26    "ZAI_API_KEY",
27    "OPENROUTER_API_KEY",
28    "DEEPSEEK_API_KEY",
29    "KIMI_API_KEY",
30    "KILO_API_KEY",
31    "NOVITA_API_KEY",
32    "MINIMAX_API_KEY",
33    "MOONSHOT_API_KEY",
34    "XAI_MANAGEMENT_KEY",
35    "ANTHROPIC_ADMIN_KEY",
36    "XAI_API_KEY",
37    "GROK_API_KEY",
38    "OPENCODE_GO_API_KEY",
39    "COMMANDCODE_API_KEY",
40    "GITHUB_COPILOT_TOKEN",
41    "GH_TOKEN",
42    "GITHUB_TOKEN",
43];
44
45pub(crate) fn vendor_secret_env_vars_to_remove(keep: &[&str]) -> Vec<&'static str> {
46    VENDOR_SECRET_ENV_VARS
47        .iter()
48        .copied()
49        .filter(|var| !keep.contains(var))
50        .collect()
51}
52
53/// Follow ordinary vendor redirects without forwarding non-standard API-key
54/// headers to a different origin. Reqwest strips `Authorization` on sensitive
55/// redirects, but vendors also use headers such as `x-api-key`, which are not
56/// covered by that built-in list.
57pub fn same_origin_redirect_policy() -> reqwest::redirect::Policy {
58    reqwest::redirect::Policy::custom(|attempt| {
59        if attempt.previous().len() >= 10 {
60            return attempt.error("too many redirects");
61        }
62        let Some(origin) = attempt.previous().first() else {
63            return attempt.stop();
64        };
65        let target = attempt.url();
66        if target.scheme() == origin.scheme()
67            && target.host_str() == origin.host_str()
68            && target.port_or_known_default() == origin.port_or_known_default()
69        {
70            attempt.follow()
71        } else {
72            attempt.stop()
73        }
74    })
75}
76
77/// Read a response body with an upper bound.
78///
79/// Every vendor buffered the whole body with `resp.bytes()` *before* anything
80/// validated it. The widget is re-executed by Waybar every 60s, so an endpoint
81/// answering with an unbounded stream had a free hand at the machine's memory.
82/// `Content-Length` is checked first when present, then the body is read in
83/// chunks so a lying or absent length cannot get past the cap either.
84pub async fn read_body_capped(
85    mut resp: reqwest::Response,
86    max: usize,
87) -> crate::error::Result<Vec<u8>> {
88    let too_big = |n: u64| {
89        crate::error::AppError::Schema(format!(
90            "response body exceeds the {max}-byte limit ({n} bytes); refusing to buffer it"
91        ))
92    };
93    if let Some(len) = resp.content_length()
94        && len > max as u64
95    {
96        return Err(too_big(len));
97    }
98    let mut buf: Vec<u8> = Vec::new();
99    while let Some(chunk) = resp.chunk().await? {
100        if chunk.len() > max.saturating_sub(buf.len()) {
101            return Err(too_big(buf.len().saturating_add(chunk.len()) as u64));
102        }
103        buf.extend_from_slice(&chunk);
104    }
105    Ok(buf)
106}
107
108/// Stable enum used by `--vendor` and in config files.
109#[derive(
110    Debug, Clone, Copy, ValueEnum, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize,
111)]
112#[serde(rename_all = "lowercase")]
113pub enum VendorId {
114    Anthropic,
115    #[serde(rename = "anthropic_api")]
116    AnthropicApi,
117    Openai,
118    Copilot,
119    Zai,
120    Openrouter,
121    Deepseek,
122    Kimi,
123    Kilo,
124    Novita,
125    Moonshot,
126    Grok,
127    Supergrok,
128    Antigravity,
129    Cursor,
130    Minimax,
131    Kiro,
132    #[serde(rename = "nous")]
133    NousResearch,
134    #[serde(rename = "opencode-go")]
135    OpenCodeGo,
136    #[serde(rename = "commandcode")]
137    CommandCode,
138}
139
140/// How a provider authenticates. Drives what a frontend offers a provider that
141/// is not usable yet: a command to run, a variable to set, or an app to sign
142/// in to.
143#[derive(Debug, Clone, Copy, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
144#[serde(rename_all = "lowercase")]
145pub enum AuthKind {
146    /// An interactive login writes a credential file. `login_command` runs it.
147    Oauth,
148    /// An API key, from the environment or an inline `api_key` in config.
149    ApiKey,
150    /// No credential of its own — a local product's session or state file is
151    /// the login, and there is nothing for the user to paste.
152    Local,
153}
154
155impl AuthKind {
156    pub const fn as_str(self) -> &'static str {
157        match self {
158            AuthKind::Oauth => "oauth",
159            AuthKind::ApiKey => "apikey",
160            AuthKind::Local => "local",
161        }
162    }
163}
164
165impl VendorId {
166    pub fn slug(self) -> &'static str {
167        match self {
168            VendorId::Anthropic => "anthropic",
169            VendorId::AnthropicApi => "anthropic_api",
170            VendorId::Openai => "openai",
171            VendorId::Copilot => "copilot",
172            VendorId::Zai => "zai",
173            VendorId::Openrouter => "openrouter",
174            VendorId::Deepseek => "deepseek",
175            VendorId::Kimi => "kimi",
176            VendorId::Kilo => "kilo",
177            VendorId::Novita => "novita",
178            VendorId::Moonshot => "moonshot",
179            VendorId::Grok => "grok",
180            VendorId::Supergrok => "supergrok",
181            VendorId::Antigravity => "antigravity",
182            VendorId::Cursor => "cursor",
183            VendorId::Minimax => "minimax",
184            VendorId::Kiro => "kiro",
185            VendorId::NousResearch => "nous",
186            VendorId::OpenCodeGo => "opencode-go",
187            VendorId::CommandCode => "commandcode",
188        }
189    }
190
191    /// Canonical human-readable name for shared reports and compact UI labels.
192    /// Platform frontends may add context (for example, "GLM (Z.AI)" in a
193    /// wide TUI tab), but should not carry their own full vendor-name table.
194    pub fn display_name(self) -> &'static str {
195        match self {
196            VendorId::Anthropic => "Claude",
197            VendorId::AnthropicApi => "Anthropic API",
198            VendorId::Openai => "Codex",
199            VendorId::Copilot => "GitHub Copilot",
200            VendorId::Zai => "Z.AI",
201            VendorId::Openrouter => "OpenRouter",
202            VendorId::Deepseek => "DeepSeek",
203            VendorId::Kimi => "Kimi",
204            VendorId::Kilo => "Kilo",
205            VendorId::Novita => "Novita",
206            VendorId::Moonshot => "Moonshot",
207            VendorId::Grok => "Grok",
208            VendorId::Supergrok => "SuperGrok",
209            VendorId::Antigravity => "Antigravity",
210            VendorId::Cursor => "Cursor",
211            VendorId::Minimax => "MiniMax",
212            VendorId::Kiro => "Kiro",
213            VendorId::NousResearch => "Nous Research",
214            VendorId::OpenCodeGo => "OpenCode Go",
215            VendorId::CommandCode => "Command Code",
216        }
217    }
218
219    /// Glyph for a compact bar chip. Same role as [`Self::short_name`]: the
220    /// Omarchy top bar (and any other frontend) takes it from `usage --json`
221    /// rather than keeping its own provider-icon table.
222    pub const fn bar_icon(self) -> &'static str {
223        match self {
224            VendorId::Anthropic => "󰚩",
225            VendorId::AnthropicApi => "󰢗",
226            VendorId::Openai => "󱢆",
227            VendorId::Copilot => "󰊤",
228            VendorId::Zai => VendorId::Zai.short_name(),
229            VendorId::Openrouter => "󱙺",
230            VendorId::Deepseek => "󰧑",
231            VendorId::Kimi => VendorId::Kimi.short_name(),
232            VendorId::Kilo => "󰭟",
233            VendorId::Novita => "󰄔",
234            VendorId::Moonshot => VendorId::Moonshot.short_name(),
235            VendorId::Grok | VendorId::Supergrok => "󰇷",
236            VendorId::Antigravity => VendorId::Antigravity.short_name(),
237            VendorId::Cursor => "❯",
238            VendorId::Minimax => VendorId::Minimax.short_name(),
239            VendorId::Kiro => "◆",
240            VendorId::NousResearch => VendorId::NousResearch.short_name(),
241            VendorId::OpenCodeGo => VendorId::OpenCodeGo.short_name(),
242            VendorId::CommandCode => VendorId::CommandCode.short_name(),
243        }
244    }
245
246    /// Compact three-letter code for the bar. This is the single source for
247    /// `{vendor_short}` in every renderer, the `usage --json` `short_name`
248    /// field, and any frontend that wants a Waybar-style provider tag; a
249    /// second copy in a placeholder map or a QML file is how the table forks.
250    pub const fn short_name(self) -> &'static str {
251        match self {
252            VendorId::Anthropic => "cld",
253            VendorId::AnthropicApi => "aac",
254            VendorId::Openai => "gpt",
255            VendorId::Copilot => "ghc",
256            VendorId::Zai => "zai",
257            VendorId::Openrouter => "opr",
258            VendorId::Deepseek => "dsk",
259            VendorId::Kimi => "kmi",
260            VendorId::Kilo => "klo",
261            VendorId::Novita => "nvt",
262            VendorId::Moonshot => "msh",
263            VendorId::Grok => "grk",
264            VendorId::Supergrok => "sgk",
265            VendorId::Antigravity => "agy",
266            VendorId::Cursor => "cur",
267            VendorId::Minimax => "mmx",
268            VendorId::Kiro => "kir",
269            VendorId::NousResearch => "nrs",
270            VendorId::OpenCodeGo => "ocg",
271            VendorId::CommandCode => "cmc",
272        }
273    }
274
275    /// How a provider proves who you are. This is the fact a frontend needs to
276    /// say what an unconfigured provider is still missing, and it is the one
277    /// thing neither `usage --json` nor the config file carries: the report
278    /// lists only *enabled* providers, so the switched-off and the
279    /// never-credentialed are exactly the rows it cannot describe.
280    pub const fn auth_kind(self) -> AuthKind {
281        match self {
282            VendorId::Anthropic
283            | VendorId::Openai
284            | VendorId::Copilot
285            | VendorId::NousResearch
286            | VendorId::CommandCode => AuthKind::Oauth,
287            VendorId::AnthropicApi
288            | VendorId::Zai
289            | VendorId::Openrouter
290            | VendorId::Deepseek
291            | VendorId::Kimi
292            | VendorId::Kilo
293            | VendorId::Novita
294            | VendorId::Moonshot
295            | VendorId::Grok
296            | VendorId::Minimax
297            | VendorId::OpenCodeGo => AuthKind::ApiKey,
298            // No credential of their own: another local product's session is
299            // the login. Antigravity has no credential file at all (the binary
300            // probes whichever local server answers), Cursor and Kiro read the
301            // IDE's and kiro-cli's own state, and SuperGrok uses the Grok Build
302            // CLI's login.
303            VendorId::Supergrok | VendorId::Antigravity | VendorId::Cursor | VendorId::Kiro => {
304                AuthKind::Local
305            }
306        }
307    }
308
309    /// Default environment variable holding this provider's key, or `""` for a
310    /// provider that has none. This is only the *default*: most key vendors
311    /// accept an `api_key_env` override in config, so a frontend showing the
312    /// variable a user must set wants [`Config::api_key_env_for`], not this.
313    pub const fn api_key_env(self) -> &'static str {
314        match self {
315            VendorId::AnthropicApi => "ANTHROPIC_ADMIN_KEY",
316            VendorId::Zai => "ZAI_API_KEY",
317            VendorId::Openrouter => "OPENROUTER_API_KEY",
318            VendorId::Deepseek => "DEEPSEEK_API_KEY",
319            VendorId::Kimi => "KIMI_API_KEY",
320            VendorId::Kilo => "KILO_API_KEY",
321            VendorId::Novita => "NOVITA_API_KEY",
322            VendorId::Moonshot => "MOONSHOT_API_KEY",
323            VendorId::Grok => "XAI_MANAGEMENT_KEY",
324            VendorId::Minimax => "MINIMAX_API_KEY",
325            VendorId::OpenCodeGo => "OPENCODE_GO_API_KEY",
326            // OAuth-first, with an environment override for CI and headless
327            // use. Neither name is configurable, so neither has an
328            // `api_key_env` field in its config section.
329            VendorId::Copilot => "GITHUB_COPILOT_TOKEN",
330            VendorId::CommandCode => "COMMANDCODE_API_KEY",
331            VendorId::Anthropic
332            | VendorId::Openai
333            | VendorId::Supergrok
334            | VendorId::Antigravity
335            | VendorId::Cursor
336            | VendorId::Kiro
337            | VendorId::NousResearch => "",
338        }
339    }
340
341    /// Command that signs this provider in, or `""` when signing in happens
342    /// somewhere this cannot name — a desktop app's own window. The strings
343    /// are the ones the vendor modules' own credential errors already print,
344    /// so a status row and a failed fetch tell the user to run the same thing.
345    pub const fn login_command(self) -> &'static str {
346        match self {
347            VendorId::Anthropic => "claude",
348            VendorId::Openai => "codex login",
349            VendorId::Copilot => "gh auth login",
350            VendorId::CommandCode => "commandcode",
351            VendorId::NousResearch => "ai-usagebar auth nous login",
352            VendorId::Kiro => "kiro-cli login",
353            // Kimi takes a key *or* the Kimi Code CLI's own OAuth login, which
354            // is what a subscriber already has locally.
355            VendorId::Kimi => "kimi",
356            VendorId::AnthropicApi
357            | VendorId::Zai
358            | VendorId::Openrouter
359            | VendorId::Deepseek
360            | VendorId::Kilo
361            | VendorId::Novita
362            | VendorId::Moonshot
363            | VendorId::Grok
364            | VendorId::Supergrok
365            | VendorId::Antigravity
366            | VendorId::Cursor
367            | VendorId::Minimax
368            | VendorId::OpenCodeGo => "",
369        }
370    }
371
372    pub fn all() -> &'static [VendorId] {
373        &[
374            VendorId::Anthropic,
375            VendorId::AnthropicApi,
376            VendorId::Openai,
377            VendorId::Copilot,
378            VendorId::Zai,
379            VendorId::Openrouter,
380            VendorId::Deepseek,
381            VendorId::Kimi,
382            VendorId::Kilo,
383            VendorId::Novita,
384            VendorId::Moonshot,
385            VendorId::Grok,
386            VendorId::Supergrok,
387            VendorId::Antigravity,
388            VendorId::Cursor,
389            VendorId::Minimax,
390            VendorId::Kiro,
391            VendorId::NousResearch,
392            VendorId::OpenCodeGo,
393            VendorId::CommandCode,
394        ]
395    }
396}
397
398/// What a vendor returns from a successful fetch — the same
399/// [`Outcome`](crate::outcome::Outcome) every vendor produces, once its own
400/// snapshot type has been widened to [`VendorSnapshot`]. Each vendor gets
401/// there with a single `outcome.map(VendorSnapshot::Whichever)`.
402pub type VendorOutcome = crate::outcome::Outcome<VendorSnapshot>;
403
404/// Options forwarded to renderers from the CLI.
405#[derive(Debug, Clone)]
406pub struct RenderOpts {
407    pub format: Option<String>,
408    pub tooltip_format: Option<String>,
409    pub icon: Option<String>,
410    pub pace_tolerance: u32,
411    pub format_pace_color: bool,
412    pub tooltip_pace_pts: bool,
413}
414
415impl RenderOpts {
416    pub fn from_cli(cli: &Cli) -> Self {
417        Self {
418            format: cli.format.clone(),
419            tooltip_format: cli.tooltip_format.clone(),
420            icon: cli.icon.clone(),
421            pace_tolerance: cli.pace_tolerance,
422            format_pace_color: cli.format_pace_color,
423            tooltip_pace_pts: cli.tooltip_pace_pts,
424        }
425    }
426}
427
428#[cfg(test)]
429mod tests {
430    use super::*;
431
432    #[test]
433    fn every_vendor_has_stable_machine_and_display_names() {
434        for vendor in VendorId::all() {
435            assert!(!vendor.slug().is_empty());
436            assert!(!vendor.display_name().is_empty());
437        }
438        assert_eq!(VendorId::Anthropic.slug(), "anthropic");
439        assert_eq!(VendorId::Anthropic.display_name(), "Claude");
440        assert_eq!(VendorId::Openai.display_name(), "Codex");
441        assert_eq!(VendorId::Zai.display_name(), "Z.AI");
442    }
443
444    /// `{vendor_short}` is a documented format placeholder and now also rides
445    /// the `usage --json` report, so a duplicate or a re-typed code would make
446    /// two providers indistinguishable in a bar that shows nothing else.
447    #[test]
448    fn every_vendor_short_name_is_a_unique_three_letter_code() {
449        let mut seen = std::collections::BTreeSet::new();
450        for vendor in VendorId::all() {
451            let short = vendor.short_name();
452            assert_eq!(short.len(), 3, "{} is not three letters", vendor.slug());
453            assert!(
454                short.chars().all(|c| c.is_ascii_lowercase()),
455                "{} is not lowercase ascii",
456                vendor.slug()
457            );
458            assert!(seen.insert(short), "{short} is used by two vendors");
459        }
460        assert_eq!(VendorId::Anthropic.short_name(), "cld");
461        assert_eq!(VendorId::Openai.short_name(), "gpt");
462        assert_eq!(VendorId::Zai.short_name(), "zai");
463        assert_eq!(VendorId::Antigravity.short_name(), "agy");
464    }
465
466    /// The bar can show every provider at once, so a glyph two providers share
467    /// tells the user nothing about which row is which. Grok and SuperGrok are
468    /// the one sanctioned pair — same brand, two products. Providers without a
469    /// distinct Nerd Font mark use their `short_name`, which is unique by
470    /// construction and cannot render as tofu.
471    #[test]
472    fn every_vendor_has_a_bar_icon_and_no_two_share_one() {
473        use std::collections::BTreeMap;
474
475        let mut by_icon: BTreeMap<&str, Vec<&str>> = BTreeMap::new();
476        for vendor in VendorId::all() {
477            assert!(!vendor.bar_icon().is_empty(), "{}", vendor.slug());
478            by_icon
479                .entry(vendor.bar_icon())
480                .or_default()
481                .push(vendor.slug());
482        }
483
484        let shared: Vec<_> = by_icon
485            .iter()
486            .filter(|(_, vendors)| vendors.len() > 1)
487            .filter(|(_, vendors)| vendors.as_slice() != ["grok", "supergrok"])
488            .collect();
489        assert!(
490            shared.is_empty(),
491            "these providers are indistinguishable in a bar that shows them \
492             side by side: {shared:#?}"
493        );
494        assert_eq!(VendorId::Anthropic.bar_icon(), "󰚩");
495        assert_eq!(VendorId::Openai.bar_icon(), "󱢆");
496        assert_eq!(VendorId::Supergrok.bar_icon(), VendorId::Grok.bar_icon());
497        assert_eq!(VendorId::CommandCode.bar_icon(), "cmc");
498    }
499
500    #[test]
501    fn new_vendor_contracts_keep_public_names_and_slugs() {
502        assert_eq!(VendorId::NousResearch.slug(), "nous");
503        assert_eq!(VendorId::NousResearch.display_name(), "Nous Research");
504        assert_eq!(VendorId::OpenCodeGo.slug(), "opencode-go");
505        assert_eq!(VendorId::OpenCodeGo.display_name(), "OpenCode Go");
506        assert_eq!(
507            serde_json::to_value(VendorId::OpenCodeGo).unwrap(),
508            serde_json::json!("opencode-go")
509        );
510    }
511
512    #[test]
513    fn vendor_secret_env_vars_cover_config_defaults() {
514        let configured_defaults = [
515            "ZAI_API_KEY",
516            "OPENROUTER_API_KEY",
517            "DEEPSEEK_API_KEY",
518            "KIMI_API_KEY",
519            "KILO_API_KEY",
520            "NOVITA_API_KEY",
521            "MINIMAX_API_KEY",
522            "MOONSHOT_API_KEY",
523            "XAI_MANAGEMENT_KEY",
524            "ANTHROPIC_ADMIN_KEY",
525            "GITHUB_COPILOT_TOKEN",
526        ];
527        for name in configured_defaults {
528            assert!(VENDOR_SECRET_ENV_VARS.contains(&name), "missing {name}");
529        }
530    }
531
532    #[test]
533    fn vars_to_remove_preserves_only_requested_grok_credentials() {
534        let removed = vendor_secret_env_vars_to_remove(&["XAI_API_KEY", "GROK_API_KEY"]);
535        assert!(!removed.contains(&"XAI_API_KEY"));
536        assert!(!removed.contains(&"GROK_API_KEY"));
537        assert!(removed.contains(&"ANTHROPIC_ADMIN_KEY"));
538        assert!(removed.contains(&"OPENROUTER_API_KEY"));
539        assert_eq!(removed.len(), VENDOR_SECRET_ENV_VARS.len() - 2);
540    }
541
542    #[test]
543    fn copilot_token_is_removed_before_unrelated_subprocesses_launch() {
544        let removed = vendor_secret_env_vars_to_remove(&[]);
545        assert!(removed.contains(&"GITHUB_COPILOT_TOKEN"));
546    }
547
548    #[tokio::test]
549    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
550        let mut server = mockito::Server::new_async().await;
551        server
552            .mock("GET", "/big")
553            .with_status(200)
554            .with_body("x".repeat(4096))
555            .create_async()
556            .await;
557        server
558            .mock("GET", "/small")
559            .with_status(200)
560            .with_body("hello")
561            .create_async()
562            .await;
563
564        let client = reqwest::Client::new();
565
566        // Over the cap: refused rather than buffered.
567        let resp = client
568            .get(format!("{}/big", server.url()))
569            .send()
570            .await
571            .unwrap();
572        let err = read_body_capped(resp, 1024).await.unwrap_err();
573        assert!(
574            err.to_string().contains("exceeds"),
575            "unexpected error: {err}"
576        );
577
578        // Under the cap: identical to the previous `resp.bytes()` behaviour.
579        let resp = client
580            .get(format!("{}/small", server.url()))
581            .send()
582            .await
583            .unwrap();
584        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
585    }
586
587    #[tokio::test]
588    async fn chunked_body_without_content_length_still_hits_the_cap() {
589        let mut server = mockito::Server::new_async().await;
590        server
591            .mock("GET", "/chunked")
592            .with_status(200)
593            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
594            .create_async()
595            .await;
596
597        let response = reqwest::Client::new()
598            .get(format!("{}/chunked", server.url()))
599            .send()
600            .await
601            .unwrap();
602        assert!(response.content_length().is_none());
603        let error = read_body_capped(response, 1024).await.unwrap_err();
604        assert!(error.to_string().contains("exceeds"), "{error}");
605    }
606
607    #[tokio::test]
608    async fn same_origin_redirects_still_work_with_vendor_headers() {
609        let mut server = mockito::Server::new_async().await;
610        let redirect = server
611            .mock("GET", "/start")
612            .match_header("x-api-key", "secret")
613            .with_status(302)
614            .with_header("location", "/finish")
615            .create_async()
616            .await;
617        let finish = server
618            .mock("GET", "/finish")
619            .match_header("x-api-key", "secret")
620            .with_status(200)
621            .create_async()
622            .await;
623        let client = reqwest::Client::builder()
624            .redirect(same_origin_redirect_policy())
625            .build()
626            .unwrap();
627
628        let response = client
629            .get(format!("{}/start", server.url()))
630            .header("x-api-key", "secret")
631            .send()
632            .await
633            .unwrap();
634
635        assert_eq!(response.status(), reqwest::StatusCode::OK);
636        redirect.assert_async().await;
637        finish.assert_async().await;
638    }
639
640    #[tokio::test]
641    async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
642        let mut origin = mockito::Server::new_async().await;
643        let mut target = mockito::Server::new_async().await;
644        let target_url = format!("{}/capture", target.url());
645        let redirect = origin
646            .mock("GET", "/start")
647            .match_header("x-api-key", "secret")
648            .with_status(302)
649            .with_header("location", &target_url)
650            .create_async()
651            .await;
652        let capture = target
653            .mock("GET", "/capture")
654            .expect(0)
655            .create_async()
656            .await;
657        let client = reqwest::Client::builder()
658            .redirect(same_origin_redirect_policy())
659            .build()
660            .unwrap();
661
662        let response = client
663            .get(format!("{}/start", origin.url()))
664            .header("x-api-key", "secret")
665            .send()
666            .await
667            .unwrap();
668
669        assert_eq!(response.status(), reqwest::StatusCode::FOUND);
670        redirect.assert_async().await;
671        capture.assert_async().await;
672    }
673
674    #[test]
675    fn vendor_id_slug_round_trip() {
676        for id in VendorId::all() {
677            assert_eq!(
678                id.slug(),
679                serde_json::to_value(id).unwrap().as_str().unwrap()
680            );
681        }
682    }
683}
ai_usagebar/vendor.rs

ai_usagebar/
vendor.rs
