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/// Follow ordinary vendor redirects without forwarding non-standard API-key
24/// headers to a different origin. Reqwest strips `Authorization` on sensitive
25/// redirects, but vendors also use headers such as `x-api-key`, which are not
26/// covered by that built-in list.
27pub fn same_origin_redirect_policy() -> reqwest::redirect::Policy {
28    reqwest::redirect::Policy::custom(|attempt| {
29        if attempt.previous().len() >= 10 {
30            return attempt.error("too many redirects");
31        }
32        let Some(origin) = attempt.previous().first() else {
33            return attempt.stop();
34        };
35        let target = attempt.url();
36        if target.scheme() == origin.scheme()
37            && target.host_str() == origin.host_str()
38            && target.port_or_known_default() == origin.port_or_known_default()
39        {
40            attempt.follow()
41        } else {
42            attempt.stop()
43        }
44    })
45}
46
47/// Read a response body with an upper bound.
48///
49/// Every vendor buffered the whole body with `resp.bytes()` *before* anything
50/// validated it. The widget is re-executed by Waybar every 60s, so an endpoint
51/// answering with an unbounded stream had a free hand at the machine's memory.
52/// `Content-Length` is checked first when present, then the body is read in
53/// chunks so a lying or absent length cannot get past the cap either.
54pub async fn read_body_capped(
55    mut resp: reqwest::Response,
56    max: usize,
57) -> crate::error::Result<Vec<u8>> {
58    let too_big = |n: u64| {
59        crate::error::AppError::Schema(format!(
60            "response body exceeds the {max}-byte limit ({n} bytes); refusing to buffer it"
61        ))
62    };
63    if let Some(len) = resp.content_length()
64        && len > max as u64
65    {
66        return Err(too_big(len));
67    }
68    let mut buf: Vec<u8> = Vec::new();
69    while let Some(chunk) = resp.chunk().await? {
70        if chunk.len() > max.saturating_sub(buf.len()) {
71            return Err(too_big(buf.len().saturating_add(chunk.len()) as u64));
72        }
73        buf.extend_from_slice(&chunk);
74    }
75    Ok(buf)
76}
77
78/// Stable enum used by `--vendor` and in config files.
79#[derive(
80    Debug, Clone, Copy, ValueEnum, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize,
81)]
82#[serde(rename_all = "lowercase")]
83pub enum VendorId {
84    Anthropic,
85    #[serde(rename = "anthropic_api")]
86    AnthropicApi,
87    Openai,
88    Zai,
89    Openrouter,
90    Deepseek,
91    Kimi,
92    Kilo,
93    Novita,
94    Moonshot,
95    Grok,
96    Supergrok,
97    Antigravity,
98    Cursor,
99    Minimax,
100    Kiro,
101}
102
103impl VendorId {
104    pub fn slug(self) -> &'static str {
105        match self {
106            VendorId::Anthropic => "anthropic",
107            VendorId::AnthropicApi => "anthropic_api",
108            VendorId::Openai => "openai",
109            VendorId::Zai => "zai",
110            VendorId::Openrouter => "openrouter",
111            VendorId::Deepseek => "deepseek",
112            VendorId::Kimi => "kimi",
113            VendorId::Kilo => "kilo",
114            VendorId::Novita => "novita",
115            VendorId::Moonshot => "moonshot",
116            VendorId::Grok => "grok",
117            VendorId::Supergrok => "supergrok",
118            VendorId::Antigravity => "antigravity",
119            VendorId::Cursor => "cursor",
120            VendorId::Minimax => "minimax",
121            VendorId::Kiro => "kiro",
122        }
123    }
124
125    /// Canonical human-readable name for shared reports and compact UI labels.
126    /// Platform frontends may add context (for example, "GLM (Z.AI)" in a
127    /// wide TUI tab), but should not carry their own full vendor-name table.
128    pub fn display_name(self) -> &'static str {
129        match self {
130            VendorId::Anthropic => "Claude",
131            VendorId::AnthropicApi => "Anthropic API",
132            VendorId::Openai => "Codex",
133            VendorId::Zai => "Z.AI",
134            VendorId::Openrouter => "OpenRouter",
135            VendorId::Deepseek => "DeepSeek",
136            VendorId::Kimi => "Kimi",
137            VendorId::Kilo => "Kilo",
138            VendorId::Novita => "Novita",
139            VendorId::Moonshot => "Moonshot",
140            VendorId::Grok => "Grok",
141            VendorId::Supergrok => "SuperGrok",
142            VendorId::Antigravity => "Antigravity",
143            VendorId::Cursor => "Cursor",
144            VendorId::Minimax => "MiniMax",
145            VendorId::Kiro => "Kiro",
146        }
147    }
148
149    pub fn all() -> &'static [VendorId] {
150        &[
151            VendorId::Anthropic,
152            VendorId::AnthropicApi,
153            VendorId::Openai,
154            VendorId::Zai,
155            VendorId::Openrouter,
156            VendorId::Deepseek,
157            VendorId::Kimi,
158            VendorId::Kilo,
159            VendorId::Novita,
160            VendorId::Moonshot,
161            VendorId::Grok,
162            VendorId::Supergrok,
163            VendorId::Antigravity,
164            VendorId::Cursor,
165            VendorId::Minimax,
166            VendorId::Kiro,
167        ]
168    }
169}
170
171/// What a vendor returns from a successful fetch — snapshot + meta. Mirrors
172/// `anthropic::fetch::FetchOutcome` but vendor-agnostic.
173#[derive(Debug, Clone)]
174pub struct VendorOutcome {
175    pub snapshot: VendorSnapshot,
176    pub stale: bool,
177    pub last_error: Option<(u16, String)>,
178    pub cache_age: Option<std::time::Duration>,
179}
180
181/// Options forwarded to renderers from the CLI.
182#[derive(Debug, Clone)]
183pub struct RenderOpts {
184    pub format: Option<String>,
185    pub tooltip_format: Option<String>,
186    pub icon: Option<String>,
187    pub pace_tolerance: u32,
188    pub format_pace_color: bool,
189    pub tooltip_pace_pts: bool,
190}
191
192impl RenderOpts {
193    pub fn from_cli(cli: &Cli) -> Self {
194        Self {
195            format: cli.format.clone(),
196            tooltip_format: cli.tooltip_format.clone(),
197            icon: cli.icon.clone(),
198            pace_tolerance: cli.pace_tolerance,
199            format_pace_color: cli.format_pace_color,
200            tooltip_pace_pts: cli.tooltip_pace_pts,
201        }
202    }
203}
204
205#[cfg(test)]
206mod tests {
207    use super::*;
208
209    #[test]
210    fn every_vendor_has_stable_machine_and_display_names() {
211        for vendor in VendorId::all() {
212            assert!(!vendor.slug().is_empty());
213            assert!(!vendor.display_name().is_empty());
214        }
215        assert_eq!(VendorId::Anthropic.slug(), "anthropic");
216        assert_eq!(VendorId::Anthropic.display_name(), "Claude");
217        assert_eq!(VendorId::Openai.display_name(), "Codex");
218        assert_eq!(VendorId::Zai.display_name(), "Z.AI");
219    }
220
221    #[tokio::test]
222    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
223        let mut server = mockito::Server::new_async().await;
224        server
225            .mock("GET", "/big")
226            .with_status(200)
227            .with_body("x".repeat(4096))
228            .create_async()
229            .await;
230        server
231            .mock("GET", "/small")
232            .with_status(200)
233            .with_body("hello")
234            .create_async()
235            .await;
236
237        let client = reqwest::Client::new();
238
239        // Over the cap: refused rather than buffered.
240        let resp = client
241            .get(format!("{}/big", server.url()))
242            .send()
243            .await
244            .unwrap();
245        let err = read_body_capped(resp, 1024).await.unwrap_err();
246        assert!(
247            err.to_string().contains("exceeds"),
248            "unexpected error: {err}"
249        );
250
251        // Under the cap: identical to the previous `resp.bytes()` behaviour.
252        let resp = client
253            .get(format!("{}/small", server.url()))
254            .send()
255            .await
256            .unwrap();
257        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
258    }
259
260    #[tokio::test]
261    async fn chunked_body_without_content_length_still_hits_the_cap() {
262        let mut server = mockito::Server::new_async().await;
263        server
264            .mock("GET", "/chunked")
265            .with_status(200)
266            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
267            .create_async()
268            .await;
269
270        let response = reqwest::Client::new()
271            .get(format!("{}/chunked", server.url()))
272            .send()
273            .await
274            .unwrap();
275        assert!(response.content_length().is_none());
276        let error = read_body_capped(response, 1024).await.unwrap_err();
277        assert!(error.to_string().contains("exceeds"), "{error}");
278    }
279
280    #[tokio::test]
281    async fn same_origin_redirects_still_work_with_vendor_headers() {
282        let mut server = mockito::Server::new_async().await;
283        let redirect = server
284            .mock("GET", "/start")
285            .match_header("x-api-key", "secret")
286            .with_status(302)
287            .with_header("location", "/finish")
288            .create_async()
289            .await;
290        let finish = server
291            .mock("GET", "/finish")
292            .match_header("x-api-key", "secret")
293            .with_status(200)
294            .create_async()
295            .await;
296        let client = reqwest::Client::builder()
297            .redirect(same_origin_redirect_policy())
298            .build()
299            .unwrap();
300
301        let response = client
302            .get(format!("{}/start", server.url()))
303            .header("x-api-key", "secret")
304            .send()
305            .await
306            .unwrap();
307
308        assert_eq!(response.status(), reqwest::StatusCode::OK);
309        redirect.assert_async().await;
310        finish.assert_async().await;
311    }
312
313    #[tokio::test]
314    async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
315        let mut origin = mockito::Server::new_async().await;
316        let mut target = mockito::Server::new_async().await;
317        let target_url = format!("{}/capture", target.url());
318        let redirect = origin
319            .mock("GET", "/start")
320            .match_header("x-api-key", "secret")
321            .with_status(302)
322            .with_header("location", &target_url)
323            .create_async()
324            .await;
325        let capture = target
326            .mock("GET", "/capture")
327            .expect(0)
328            .create_async()
329            .await;
330        let client = reqwest::Client::builder()
331            .redirect(same_origin_redirect_policy())
332            .build()
333            .unwrap();
334
335        let response = client
336            .get(format!("{}/start", origin.url()))
337            .header("x-api-key", "secret")
338            .send()
339            .await
340            .unwrap();
341
342        assert_eq!(response.status(), reqwest::StatusCode::FOUND);
343        redirect.assert_async().await;
344        capture.assert_async().await;
345    }
346
347    #[test]
348    fn vendor_id_slug_round_trip() {
349        for id in VendorId::all() {
350            assert_eq!(
351                id.slug(),
352                serde_json::to_value(id).unwrap().as_str().unwrap()
353            );
354        }
355    }
356}
ai_usagebar/vendor.rs

ai_usagebar/
vendor.rs
