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/// Read a response body with an upper bound.
24///
25/// Every vendor buffered the whole body with `resp.bytes()` *before* anything
26/// validated it. The widget is re-executed by Waybar every 60s, so an endpoint
27/// answering with an unbounded stream had a free hand at the machine's memory.
28/// `Content-Length` is checked first when present, then the body is read in
29/// chunks so a lying or absent length cannot get past the cap either.
30pub async fn read_body_capped(
31    mut resp: reqwest::Response,
32    max: usize,
33) -> crate::error::Result<Vec<u8>> {
34    let too_big = |n: u64| {
35        crate::error::AppError::Schema(format!(
36            "response body exceeds the {max}-byte limit ({n} bytes); refusing to buffer it"
37        ))
38    };
39    if let Some(len) = resp.content_length()
40        && len > max as u64
41    {
42        return Err(too_big(len));
43    }
44    let mut buf: Vec<u8> = Vec::new();
45    while let Some(chunk) = resp.chunk().await? {
46        if chunk.len() > max.saturating_sub(buf.len()) {
47            return Err(too_big(buf.len().saturating_add(chunk.len()) as u64));
48        }
49        buf.extend_from_slice(&chunk);
50    }
51    Ok(buf)
52}
53
54/// Stable enum used by `--vendor` and in config files.
55#[derive(
56    Debug, Clone, Copy, ValueEnum, PartialEq, Eq, Hash, serde::Deserialize, serde::Serialize,
57)]
58#[serde(rename_all = "lowercase")]
59pub enum VendorId {
60    Anthropic,
61    #[serde(rename = "anthropic_api")]
62    AnthropicApi,
63    Openai,
64    Zai,
65    Openrouter,
66    Deepseek,
67    Kimi,
68    Kilo,
69    Novita,
70    Moonshot,
71    Grok,
72    Antigravity,
73    Cursor,
74    Minimax,
75}
76
77impl VendorId {
78    pub fn slug(self) -> &'static str {
79        match self {
80            VendorId::Anthropic => "anthropic",
81            VendorId::AnthropicApi => "anthropic_api",
82            VendorId::Openai => "openai",
83            VendorId::Zai => "zai",
84            VendorId::Openrouter => "openrouter",
85            VendorId::Deepseek => "deepseek",
86            VendorId::Kimi => "kimi",
87            VendorId::Kilo => "kilo",
88            VendorId::Novita => "novita",
89            VendorId::Moonshot => "moonshot",
90            VendorId::Grok => "grok",
91            VendorId::Antigravity => "antigravity",
92            VendorId::Cursor => "cursor",
93            VendorId::Minimax => "minimax",
94        }
95    }
96
97    pub fn all() -> &'static [VendorId] {
98        &[
99            VendorId::Anthropic,
100            VendorId::AnthropicApi,
101            VendorId::Openai,
102            VendorId::Zai,
103            VendorId::Openrouter,
104            VendorId::Deepseek,
105            VendorId::Kimi,
106            VendorId::Kilo,
107            VendorId::Novita,
108            VendorId::Moonshot,
109            VendorId::Grok,
110            VendorId::Antigravity,
111            VendorId::Cursor,
112            VendorId::Minimax,
113        ]
114    }
115}
116
117/// What a vendor returns from a successful fetch — snapshot + meta. Mirrors
118/// `anthropic::fetch::FetchOutcome` but vendor-agnostic.
119#[derive(Debug, Clone)]
120pub struct VendorOutcome {
121    pub snapshot: VendorSnapshot,
122    pub stale: bool,
123    pub last_error: Option<(u16, String)>,
124    pub cache_age: Option<std::time::Duration>,
125}
126
127/// Options forwarded to renderers from the CLI.
128#[derive(Debug, Clone)]
129pub struct RenderOpts {
130    pub format: Option<String>,
131    pub tooltip_format: Option<String>,
132    pub icon: Option<String>,
133    pub pace_tolerance: u32,
134    pub format_pace_color: bool,
135    pub tooltip_pace_pts: bool,
136}
137
138impl RenderOpts {
139    pub fn from_cli(cli: &Cli) -> Self {
140        Self {
141            format: cli.format.clone(),
142            tooltip_format: cli.tooltip_format.clone(),
143            icon: cli.icon.clone(),
144            pace_tolerance: cli.pace_tolerance,
145            format_pace_color: cli.format_pace_color,
146            tooltip_pace_pts: cli.tooltip_pace_pts,
147        }
148    }
149}
150
151#[cfg(test)]
152mod tests {
153    use super::*;
154
155    #[tokio::test]
156    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
157        let mut server = mockito::Server::new_async().await;
158        server
159            .mock("GET", "/big")
160            .with_status(200)
161            .with_body("x".repeat(4096))
162            .create_async()
163            .await;
164        server
165            .mock("GET", "/small")
166            .with_status(200)
167            .with_body("hello")
168            .create_async()
169            .await;
170
171        let client = reqwest::Client::new();
172
173        // Over the cap: refused rather than buffered.
174        let resp = client
175            .get(format!("{}/big", server.url()))
176            .send()
177            .await
178            .unwrap();
179        let err = read_body_capped(resp, 1024).await.unwrap_err();
180        assert!(
181            err.to_string().contains("exceeds"),
182            "unexpected error: {err}"
183        );
184
185        // Under the cap: identical to the previous `resp.bytes()` behaviour.
186        let resp = client
187            .get(format!("{}/small", server.url()))
188            .send()
189            .await
190            .unwrap();
191        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
192    }
193
194    #[tokio::test]
195    async fn chunked_body_without_content_length_still_hits_the_cap() {
196        let mut server = mockito::Server::new_async().await;
197        server
198            .mock("GET", "/chunked")
199            .with_status(200)
200            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
201            .create_async()
202            .await;
203
204        let response = reqwest::Client::new()
205            .get(format!("{}/chunked", server.url()))
206            .send()
207            .await
208            .unwrap();
209        assert!(response.content_length().is_none());
210        let error = read_body_capped(response, 1024).await.unwrap_err();
211        assert!(error.to_string().contains("exceeds"), "{error}");
212    }
213
214    #[test]
215    fn vendor_id_slug_round_trip() {
216        for id in VendorId::all() {
217            assert_eq!(
218                id.slug(),
219                serde_json::to_value(id).unwrap().as_str().unwrap()
220            );
221        }
222    }
223}
ai_usagebar/vendor.rs

ai_usagebar/
vendor.rs
