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

ai_usagebar/
vendor.rs
