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    pub fn all() -> &'static [VendorId] {
126        &[
127            VendorId::Anthropic,
128            VendorId::AnthropicApi,
129            VendorId::Openai,
130            VendorId::Zai,
131            VendorId::Openrouter,
132            VendorId::Deepseek,
133            VendorId::Kimi,
134            VendorId::Kilo,
135            VendorId::Novita,
136            VendorId::Moonshot,
137            VendorId::Grok,
138            VendorId::Supergrok,
139            VendorId::Antigravity,
140            VendorId::Cursor,
141            VendorId::Minimax,
142            VendorId::Kiro,
143        ]
144    }
145}
146
147/// What a vendor returns from a successful fetch — snapshot + meta. Mirrors
148/// `anthropic::fetch::FetchOutcome` but vendor-agnostic.
149#[derive(Debug, Clone)]
150pub struct VendorOutcome {
151    pub snapshot: VendorSnapshot,
152    pub stale: bool,
153    pub last_error: Option<(u16, String)>,
154    pub cache_age: Option<std::time::Duration>,
155}
156
157/// Options forwarded to renderers from the CLI.
158#[derive(Debug, Clone)]
159pub struct RenderOpts {
160    pub format: Option<String>,
161    pub tooltip_format: Option<String>,
162    pub icon: Option<String>,
163    pub pace_tolerance: u32,
164    pub format_pace_color: bool,
165    pub tooltip_pace_pts: bool,
166}
167
168impl RenderOpts {
169    pub fn from_cli(cli: &Cli) -> Self {
170        Self {
171            format: cli.format.clone(),
172            tooltip_format: cli.tooltip_format.clone(),
173            icon: cli.icon.clone(),
174            pace_tolerance: cli.pace_tolerance,
175            format_pace_color: cli.format_pace_color,
176            tooltip_pace_pts: cli.tooltip_pace_pts,
177        }
178    }
179}
180
181#[cfg(test)]
182mod tests {
183    use super::*;
184
185    #[tokio::test]
186    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
187        let mut server = mockito::Server::new_async().await;
188        server
189            .mock("GET", "/big")
190            .with_status(200)
191            .with_body("x".repeat(4096))
192            .create_async()
193            .await;
194        server
195            .mock("GET", "/small")
196            .with_status(200)
197            .with_body("hello")
198            .create_async()
199            .await;
200
201        let client = reqwest::Client::new();
202
203        // Over the cap: refused rather than buffered.
204        let resp = client
205            .get(format!("{}/big", server.url()))
206            .send()
207            .await
208            .unwrap();
209        let err = read_body_capped(resp, 1024).await.unwrap_err();
210        assert!(
211            err.to_string().contains("exceeds"),
212            "unexpected error: {err}"
213        );
214
215        // Under the cap: identical to the previous `resp.bytes()` behaviour.
216        let resp = client
217            .get(format!("{}/small", server.url()))
218            .send()
219            .await
220            .unwrap();
221        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
222    }
223
224    #[tokio::test]
225    async fn chunked_body_without_content_length_still_hits_the_cap() {
226        let mut server = mockito::Server::new_async().await;
227        server
228            .mock("GET", "/chunked")
229            .with_status(200)
230            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
231            .create_async()
232            .await;
233
234        let response = reqwest::Client::new()
235            .get(format!("{}/chunked", server.url()))
236            .send()
237            .await
238            .unwrap();
239        assert!(response.content_length().is_none());
240        let error = read_body_capped(response, 1024).await.unwrap_err();
241        assert!(error.to_string().contains("exceeds"), "{error}");
242    }
243
244    #[tokio::test]
245    async fn same_origin_redirects_still_work_with_vendor_headers() {
246        let mut server = mockito::Server::new_async().await;
247        let redirect = server
248            .mock("GET", "/start")
249            .match_header("x-api-key", "secret")
250            .with_status(302)
251            .with_header("location", "/finish")
252            .create_async()
253            .await;
254        let finish = server
255            .mock("GET", "/finish")
256            .match_header("x-api-key", "secret")
257            .with_status(200)
258            .create_async()
259            .await;
260        let client = reqwest::Client::builder()
261            .redirect(same_origin_redirect_policy())
262            .build()
263            .unwrap();
264
265        let response = client
266            .get(format!("{}/start", server.url()))
267            .header("x-api-key", "secret")
268            .send()
269            .await
270            .unwrap();
271
272        assert_eq!(response.status(), reqwest::StatusCode::OK);
273        redirect.assert_async().await;
274        finish.assert_async().await;
275    }
276
277    #[tokio::test]
278    async fn cross_origin_redirects_are_not_followed_with_vendor_headers() {
279        let mut origin = mockito::Server::new_async().await;
280        let mut target = mockito::Server::new_async().await;
281        let target_url = format!("{}/capture", target.url());
282        let redirect = origin
283            .mock("GET", "/start")
284            .match_header("x-api-key", "secret")
285            .with_status(302)
286            .with_header("location", &target_url)
287            .create_async()
288            .await;
289        let capture = target
290            .mock("GET", "/capture")
291            .expect(0)
292            .create_async()
293            .await;
294        let client = reqwest::Client::builder()
295            .redirect(same_origin_redirect_policy())
296            .build()
297            .unwrap();
298
299        let response = client
300            .get(format!("{}/start", origin.url()))
301            .header("x-api-key", "secret")
302            .send()
303            .await
304            .unwrap();
305
306        assert_eq!(response.status(), reqwest::StatusCode::FOUND);
307        redirect.assert_async().await;
308        capture.assert_async().await;
309    }
310
311    #[test]
312    fn vendor_id_slug_round_trip() {
313        for id in VendorId::all() {
314            assert_eq!(
315                id.slug(),
316                serde_json::to_value(id).unwrap().as_str().unwrap()
317            );
318        }
319    }
320}
ai_usagebar/vendor.rs

ai_usagebar/
vendor.rs
