Skip to main content

ai_usagebar/openai/
vendor.rs

1//! OpenAI renderer — mirrors anthropic's layout closely since both expose the
2//! same primary+secondary window pattern.
3
4use std::collections::HashMap;
5
6use chrono::{DateTime, Utc};
7
8use crate::countdown;
9use crate::format::{placeholders, substitute, updated_at_hm};
10use crate::pacing::{self, PaceSeverity};
11use crate::pango::{color_span, escape, severity_color, severity_for};
12use crate::theme::Theme;
13use crate::tooltip::{Line as TooltipLine, push_window, render_bordered};
14use crate::usage::{OpenAiSnapshot, OpenAiSource};
15use crate::vendor::{RenderOpts, VendorOutcome};
16use crate::waybar::{Class, WaybarOutput};
17
18use super::fetch::FetchOutcome;
19
20pub const DEFAULT_FORMAT: &str = "{oai_session_pct}% · {oai_session_reset}";
21
22pub fn build_placeholders(
23    snap: &OpenAiSnapshot,
24    opts: &RenderOpts,
25    now: DateTime<Utc>,
26) -> HashMap<&'static str, String> {
27    let session_p = pacing::calc(
28        snap.session.utilization_pct,
29        snap.session.resets_at,
30        now,
31        snap.session.window_duration,
32        opts.pace_tolerance,
33    );
34    let weekly_p = pacing::calc(
35        snap.weekly.utilization_pct,
36        snap.weekly.resets_at,
37        now,
38        snap.weekly.window_duration,
39        opts.pace_tolerance,
40    );
41    let cr_pct = snap
42        .code_review
43        .as_ref()
44        .map(|w| w.utilization_pct)
45        .unwrap_or(0);
46    let credit_balance = snap
47        .credits
48        .as_ref()
49        .map(|c| c.balance.clone())
50        .unwrap_or_else(|| "n/a".into());
51    let local_msgs = snap
52        .credits
53        .as_ref()
54        .and_then(|c| c.approx_local_messages)
55        .map(|(a, b)| format!("{a}-{b}"))
56        .unwrap_or_default();
57    let cloud_msgs = snap
58        .credits
59        .as_ref()
60        .and_then(|c| c.approx_cloud_messages)
61        .map(|(a, b)| format!("{a}-{b}"))
62        .unwrap_or_default();
63
64    placeholders(vec![
65        ("icon", "󱢆".to_string()),
66        ("vendor_short", "gpt".to_string()),
67        // Cross-vendor aliases — same names work across all vendors so a
68        // single `--format '{vendor_short} {session_pct}% · {session_reset}'`
69        // renders correctly during scroll-cycle.
70        ("session_pct", snap.session.utilization_pct.to_string()),
71        (
72            "session_reset",
73            countdown::format(snap.session.resets_at, now),
74        ),
75        ("weekly_pct", snap.weekly.utilization_pct.to_string()),
76        (
77            "weekly_reset",
78            countdown::format(snap.weekly.resets_at, now),
79        ),
80        ("plan", snap.plan.clone()),
81        ("oai_plan", snap.plan.clone()),
82        ("oai_session_pct", snap.session.utilization_pct.to_string()),
83        (
84            "oai_session_reset",
85            countdown::format(snap.session.resets_at, now),
86        ),
87        ("oai_session_elapsed", session_p.elapsed_pct.to_string()),
88        ("oai_session_pace", session_p.ratio_pace.glyph().to_string()),
89        (
90            "oai_session_pace_indicator",
91            session_p.point_pace.glyph().to_string(),
92        ),
93        ("oai_weekly_pct", snap.weekly.utilization_pct.to_string()),
94        (
95            "oai_weekly_reset",
96            countdown::format(snap.weekly.resets_at, now),
97        ),
98        ("oai_weekly_elapsed", weekly_p.elapsed_pct.to_string()),
99        ("oai_weekly_pace", weekly_p.ratio_pace.glyph().to_string()),
100        (
101            "oai_weekly_pace_indicator",
102            weekly_p.point_pace.glyph().to_string(),
103        ),
104        ("oai_code_review_pct", cr_pct.to_string()),
105        ("oai_credit_balance", credit_balance),
106        ("oai_local_msgs", local_msgs),
107        ("oai_cloud_msgs", cloud_msgs),
108    ])
109}
110
111pub fn severity(snap: &OpenAiSnapshot) -> PaceSeverity {
112    let mut max = snap
113        .session
114        .utilization_pct
115        .max(snap.weekly.utilization_pct);
116    if let Some(c) = &snap.code_review {
117        max = max.max(c.utilization_pct);
118    }
119    severity_for(max)
120}
121
122pub fn render(
123    outcome: &VendorOutcome,
124    snap: &OpenAiSnapshot,
125    theme: &Theme,
126    opts: &RenderOpts,
127    now: DateTime<Utc>,
128) -> WaybarOutput {
129    let class = Class::from(severity(snap));
130    let format = opts
131        .format
132        .clone()
133        .unwrap_or_else(|| DEFAULT_FORMAT.to_string());
134    let values = build_placeholders(snap, opts, now);
135
136    let mut text = substitute(&format, &values);
137    if outcome.stale {
138        text.push_str(" ⏸");
139    }
140    let wrapper_color = severity_color(severity(snap), theme).to_string();
141    let icon_prefix = match opts.icon.as_deref() {
142        Some(ic) if !ic.is_empty() => format!("{ic} "),
143        _ => String::new(),
144    };
145    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
146
147    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
148        substitute(fmt, &values)
149    } else {
150        render_tooltip(outcome, snap, theme, now)
151    };
152
153    WaybarOutput {
154        text: bar_text,
155        tooltip,
156        class,
157    }
158}
159
160fn render_tooltip(
161    outcome: &VendorOutcome,
162    snap: &OpenAiSnapshot,
163    theme: &Theme,
164    now: DateTime<Utc>,
165) -> String {
166    let blue = &theme.blue;
167    let dim = &theme.dim;
168    let fg = &theme.fg;
169
170    let mut lines: Vec<TooltipLine> = Vec::new();
171    lines.push(TooltipLine::Center(format!(
172        "<span font_weight='bold' foreground='{blue}'>{plan}</span>",
173        plan = escape(&snap.plan)
174    )));
175    lines.push(TooltipLine::Sep);
176    lines.push(TooltipLine::Body("".into()));
177
178    push_window(&mut lines, "  󰔟  Codex 5h", &snap.session, theme, now, None);
179    lines.push(TooltipLine::Body("".into()));
180    push_window(
181        &mut lines,
182        "  󰃰  Codex weekly",
183        &snap.weekly,
184        theme,
185        now,
186        None,
187    );
188
189    if let Some(cr) = snap.code_review.as_ref() {
190        lines.push(TooltipLine::Body("".into()));
191        push_window(
192            &mut lines,
193            "  󱦰  Code review (weekly)",
194            cr,
195            theme,
196            now,
197            None,
198        );
199    }
200
201    if let Some(c) = snap.credits.as_ref() {
202        lines.push(TooltipLine::Body("".into()));
203        lines.push(TooltipLine::Sep);
204        let label = if c.unlimited {
205            "unlimited".to_string()
206        } else {
207            c.balance.clone()
208        };
209        lines.push(TooltipLine::Body(format!(
210            " <span foreground='{fg}'>  󰄑  Credits</span>"
211        )));
212        lines.push(TooltipLine::Body(format!(
213            " <span foreground='{dim}'>     balance: {b}</span>",
214            b = escape(&label)
215        )));
216        if let Some((lo, hi)) = c.approx_local_messages {
217            lines.push(TooltipLine::Body(format!(
218                " <span foreground='{dim}'>     ~ {lo}-{hi} local messages</span>"
219            )));
220        }
221        if let Some((lo, hi)) = c.approx_cloud_messages {
222            lines.push(TooltipLine::Body(format!(
223                " <span foreground='{dim}'>     ~ {lo}-{hi} cloud messages</span>"
224            )));
225        }
226    }
227
228    if matches!(snap.source, OpenAiSource::Unavailable) {
229        lines.push(TooltipLine::Body("".into()));
230        lines.push(TooltipLine::Sep);
231        lines.push(TooltipLine::Body(format!(
232            " <span foreground='{dim}'>OpenAI plan usage requires Codex OAuth.</span>"
233        )));
234        lines.push(TooltipLine::Body(format!(
235            " <span foreground='{dim}'>Run `codex login` to enable.</span>"
236        )));
237    }
238
239    if let Some((code, msg)) = outcome.last_error.as_ref()
240        && *code != 0
241    {
242        let (icon, ecolor) = if *code >= 500 {
243            ("󰅚", theme.red.as_str())
244        } else {
245            ("󰀪", theme.orange.as_str())
246        };
247        lines.push(TooltipLine::Body("".into()));
248        lines.push(TooltipLine::Sep);
249        lines.push(TooltipLine::Body(format!(
250            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
251        )));
252        lines.push(TooltipLine::Body(format!(
253            "     <span foreground='{dim}'>{}</span>",
254            escape(msg)
255        )));
256    }
257
258    let updated = updated_at_hm(now, outcome.cache_age);
259    lines.push(TooltipLine::Body("".into()));
260    lines.push(TooltipLine::Sep);
261    lines.push(TooltipLine::Body(format!(
262        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
263    )));
264
265    render_bordered(&lines, theme)
266}
267
268impl From<FetchOutcome> for VendorOutcome {
269    fn from(o: FetchOutcome) -> Self {
270        Self {
271            snapshot: crate::usage::VendorSnapshot::Openai(o.snapshot),
272            stale: o.stale,
273            last_error: o.last_error,
274            cache_age: o.cache_age,
275        }
276    }
277}
278
279#[cfg(test)]
280mod tests {
281    use super::*;
282    use crate::usage::{OpenAiSnapshot, OpenAiSource, UsageWindow};
283
284    fn sample() -> OpenAiSnapshot {
285        OpenAiSnapshot {
286            plan: "ChatGPT Plus".into(),
287            session: UsageWindow {
288                utilization_pct: 1,
289                resets_at: Some(Utc::now() + chrono::Duration::hours(5)),
290                window_duration: chrono::Duration::hours(5),
291            },
292            weekly: UsageWindow {
293                utilization_pct: 0,
294                resets_at: Some(Utc::now() + chrono::Duration::days(7)),
295                window_duration: chrono::Duration::days(7),
296            },
297            code_review: None,
298            credits: None,
299            source: OpenAiSource::CodexOauth,
300        }
301    }
302
303    fn oc(s: OpenAiSnapshot) -> VendorOutcome {
304        VendorOutcome {
305            snapshot: crate::usage::VendorSnapshot::Openai(s),
306            stale: false,
307            last_error: None,
308            cache_age: Some(std::time::Duration::from_secs(15)),
309        }
310    }
311
312    fn opts() -> RenderOpts {
313        RenderOpts {
314            format: None,
315            tooltip_format: None,
316            icon: None,
317            pace_tolerance: 5,
318            format_pace_color: false,
319            tooltip_pace_pts: false,
320        }
321    }
322
323    #[test]
324    fn default_format_renders_session() {
325        let s = sample();
326        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
327        assert!(out.text.contains("1%"));
328    }
329
330    #[test]
331    fn tooltip_has_both_windows() {
332        let s = sample();
333        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
334        assert!(out.tooltip.contains("Codex 5h"));
335        assert!(out.tooltip.contains("Codex weekly"));
336        assert!(!out.tooltip.contains("Code review"));
337        assert!(!out.tooltip.contains("Credits"));
338    }
339
340    #[test]
341    fn tooltip_includes_credits_block_when_present() {
342        let mut s = sample();
343        s.credits = Some(crate::usage::OpenAiCredits {
344            balance: "$5.00".into(),
345            has_credits: true,
346            unlimited: false,
347            approx_local_messages: Some((100, 200)),
348            approx_cloud_messages: Some((30, 50)),
349        });
350        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
351        assert!(out.tooltip.contains("Credits"));
352        assert!(out.tooltip.contains("$5.00"));
353        assert!(out.tooltip.contains("100-200 local messages"));
354        assert!(out.tooltip.contains("30-50 cloud messages"));
355    }
356
357    #[test]
358    fn unavailable_source_shows_codex_login_hint() {
359        let mut s = sample();
360        s.source = OpenAiSource::Unavailable;
361        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
362        assert!(out.tooltip.contains("codex login"));
363    }
364
365    #[test]
366    fn severity_picks_worst_window() {
367        let mut s = sample();
368        s.weekly.utilization_pct = 95;
369        assert_eq!(severity(&s), PaceSeverity::Critical);
370    }
371}