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}
74
75impl VendorId {
76    pub fn slug(self) -> &'static str {
77        match self {
78            VendorId::Anthropic => "anthropic",
79            VendorId::AnthropicApi => "anthropic_api",
80            VendorId::Openai => "openai",
81            VendorId::Zai => "zai",
82            VendorId::Openrouter => "openrouter",
83            VendorId::Deepseek => "deepseek",
84            VendorId::Kimi => "kimi",
85            VendorId::Kilo => "kilo",
86            VendorId::Novita => "novita",
87            VendorId::Moonshot => "moonshot",
88            VendorId::Grok => "grok",
89            VendorId::Antigravity => "antigravity",
90        }
91    }
92
93    pub fn all() -> &'static [VendorId] {
94        &[
95            VendorId::Anthropic,
96            VendorId::AnthropicApi,
97            VendorId::Openai,
98            VendorId::Zai,
99            VendorId::Openrouter,
100            VendorId::Deepseek,
101            VendorId::Kimi,
102            VendorId::Kilo,
103            VendorId::Novita,
104            VendorId::Moonshot,
105            VendorId::Grok,
106            VendorId::Antigravity,
107        ]
108    }
109}
110
111/// What a vendor returns from a successful fetch — snapshot + meta. Mirrors
112/// `anthropic::fetch::FetchOutcome` but vendor-agnostic.
113#[derive(Debug, Clone)]
114pub struct VendorOutcome {
115    pub snapshot: VendorSnapshot,
116    pub stale: bool,
117    pub last_error: Option<(u16, String)>,
118    pub cache_age: Option<std::time::Duration>,
119}
120
121/// Options forwarded to renderers from the CLI.
122#[derive(Debug, Clone)]
123pub struct RenderOpts {
124    pub format: Option<String>,
125    pub tooltip_format: Option<String>,
126    pub icon: Option<String>,
127    pub pace_tolerance: u32,
128    pub format_pace_color: bool,
129    pub tooltip_pace_pts: bool,
130}
131
132impl RenderOpts {
133    pub fn from_cli(cli: &Cli) -> Self {
134        Self {
135            format: cli.format.clone(),
136            tooltip_format: cli.tooltip_format.clone(),
137            icon: cli.icon.clone(),
138            pace_tolerance: cli.pace_tolerance,
139            format_pace_color: cli.format_pace_color,
140            tooltip_pace_pts: cli.tooltip_pace_pts,
141        }
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[tokio::test]
150    async fn body_over_the_cap_is_refused_and_under_it_round_trips() {
151        let mut server = mockito::Server::new_async().await;
152        server
153            .mock("GET", "/big")
154            .with_status(200)
155            .with_body("x".repeat(4096))
156            .create_async()
157            .await;
158        server
159            .mock("GET", "/small")
160            .with_status(200)
161            .with_body("hello")
162            .create_async()
163            .await;
164
165        let client = reqwest::Client::new();
166
167        // Over the cap: refused rather than buffered.
168        let resp = client
169            .get(format!("{}/big", server.url()))
170            .send()
171            .await
172            .unwrap();
173        let err = read_body_capped(resp, 1024).await.unwrap_err();
174        assert!(
175            err.to_string().contains("exceeds"),
176            "unexpected error: {err}"
177        );
178
179        // Under the cap: identical to the previous `resp.bytes()` behaviour.
180        let resp = client
181            .get(format!("{}/small", server.url()))
182            .send()
183            .await
184            .unwrap();
185        assert_eq!(read_body_capped(resp, 1024).await.unwrap(), b"hello");
186    }
187
188    #[tokio::test]
189    async fn chunked_body_without_content_length_still_hits_the_cap() {
190        let mut server = mockito::Server::new_async().await;
191        server
192            .mock("GET", "/chunked")
193            .with_status(200)
194            .with_chunked_body(|writer| writer.write_all(&[b'x'; 4096]))
195            .create_async()
196            .await;
197
198        let response = reqwest::Client::new()
199            .get(format!("{}/chunked", server.url()))
200            .send()
201            .await
202            .unwrap();
203        assert!(response.content_length().is_none());
204        let error = read_body_capped(response, 1024).await.unwrap_err();
205        assert!(error.to_string().contains("exceeds"), "{error}");
206    }
207
208    #[test]
209    fn vendor_id_slug_round_trip() {
210        for id in VendorId::all() {
211            assert_eq!(
212                id.slug(),
213                serde_json::to_value(id).unwrap().as_str().unwrap()
214            );
215        }
216    }
217}
ai_usagebar/vendor.rs

ai_usagebar/
vendor.rs
