Skip to main content

ai_usagebar/openai/
vendor.rs

1//! OpenAI renderer — mirrors Anthropic's rolling-window layout while omitting
2//! any 5h or weekly window the Codex endpoint does not report.
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, UsageWindow};
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}";
21const WEEKLY_ONLY_FORMAT: &str = "{oai_weekly_pct}% · {oai_weekly_reset}";
22const NO_WINDOWS_FORMAT: &str = "{oai_plan}";
23
24pub fn build_placeholders(
25    snap: &OpenAiSnapshot,
26    opts: &RenderOpts,
27    now: DateTime<Utc>,
28) -> HashMap<&'static str, String> {
29    let session = window_placeholders(snap.session.as_ref(), opts, now);
30    let weekly = window_placeholders(snap.weekly.as_ref(), opts, now);
31    let cr_pct = snap
32        .code_review
33        .as_ref()
34        .map(|w| w.utilization_pct)
35        .unwrap_or(0);
36    let credit_balance = snap
37        .credits
38        .as_ref()
39        .map(|c| c.balance.clone())
40        .unwrap_or_else(|| "n/a".into());
41    let local_msgs = snap
42        .credits
43        .as_ref()
44        .and_then(|c| c.approx_local_messages)
45        .map(|(a, b)| format!("{a}-{b}"))
46        .unwrap_or_default();
47    let cloud_msgs = snap
48        .credits
49        .as_ref()
50        .and_then(|c| c.approx_cloud_messages)
51        .map(|(a, b)| format!("{a}-{b}"))
52        .unwrap_or_default();
53
54    placeholders(vec![
55        ("icon", "󱢆".to_string()),
56        ("vendor_short", "gpt".to_string()),
57        // Cross-vendor aliases — same names work across all vendors so a
58        // single `--format '{vendor_short} {session_pct}% · {session_reset}'`
59        // renders correctly during scroll-cycle.
60        ("session_pct", session.pct.clone()),
61        ("session_reset", session.reset.clone()),
62        ("weekly_pct", weekly.pct.clone()),
63        ("weekly_reset", weekly.reset.clone()),
64        ("session_elapsed", session.elapsed.clone()),
65        ("weekly_elapsed", weekly.elapsed.clone()),
66        ("plan", snap.plan.clone()),
67        ("oai_plan", snap.plan.clone()),
68        ("oai_session_pct", session.pct),
69        ("oai_session_reset", session.reset),
70        ("oai_session_elapsed", session.elapsed),
71        ("oai_session_pace", session.ratio_pace),
72        ("oai_session_pace_indicator", session.point_pace),
73        ("oai_weekly_pct", weekly.pct),
74        ("oai_weekly_reset", weekly.reset),
75        ("oai_weekly_elapsed", weekly.elapsed),
76        ("oai_weekly_pace", weekly.ratio_pace),
77        ("oai_weekly_pace_indicator", weekly.point_pace),
78        ("oai_code_review_pct", cr_pct.to_string()),
79        ("oai_credit_balance", credit_balance),
80        ("oai_local_msgs", local_msgs),
81        ("oai_cloud_msgs", cloud_msgs),
82    ])
83}
84
85#[derive(Default)]
86struct WindowPlaceholderValues {
87    pct: String,
88    reset: String,
89    elapsed: String,
90    ratio_pace: String,
91    point_pace: String,
92}
93
94fn window_placeholders(
95    window: Option<&UsageWindow>,
96    opts: &RenderOpts,
97    now: DateTime<Utc>,
98) -> WindowPlaceholderValues {
99    let Some(window) = window else {
100        return WindowPlaceholderValues::default();
101    };
102    let pace = pacing::calc(
103        window.utilization_pct,
104        window.resets_at,
105        now,
106        window.window_duration,
107        opts.pace_tolerance,
108    );
109    WindowPlaceholderValues {
110        pct: window.utilization_pct.to_string(),
111        reset: countdown::format(window.resets_at, now),
112        elapsed: pace.elapsed_pct.to_string(),
113        ratio_pace: pace.ratio_pace.glyph().to_string(),
114        point_pace: pace.point_pace.glyph().to_string(),
115    }
116}
117
118pub fn severity(snap: &OpenAiSnapshot) -> PaceSeverity {
119    let windows = [
120        snap.session.as_ref(),
121        snap.weekly.as_ref(),
122        snap.code_review.as_ref(),
123    ];
124    let max = windows
125        .into_iter()
126        .flatten()
127        .map(|window| window.utilization_pct)
128        .max()
129        .unwrap_or(0);
130    severity_for(max)
131}
132
133pub fn render(
134    outcome: &VendorOutcome,
135    snap: &OpenAiSnapshot,
136    theme: &Theme,
137    opts: &RenderOpts,
138    now: DateTime<Utc>,
139) -> WaybarOutput {
140    let class = Class::from(severity(snap));
141    let format = opts
142        .format
143        .clone()
144        .unwrap_or_else(|| default_format(snap).to_string());
145    let values = build_placeholders(snap, opts, now);
146
147    let mut text = substitute(&format, &values);
148    if outcome.stale {
149        text.push_str(" ⏸");
150    }
151    let wrapper_color = severity_color(severity(snap), theme).to_string();
152    let icon_prefix = match opts.icon.as_deref() {
153        Some(ic) if !ic.is_empty() => format!("{ic} "),
154        _ => String::new(),
155    };
156    let bar_text = color_span(&wrapper_color, &format!("{icon_prefix}{text}"));
157
158    let tooltip = if let Some(fmt) = opts.tooltip_format.as_deref() {
159        substitute(fmt, &values)
160    } else {
161        render_tooltip(outcome, snap, theme, now)
162    };
163
164    WaybarOutput {
165        text: bar_text,
166        tooltip,
167        class,
168    }
169}
170
171fn default_format(snap: &OpenAiSnapshot) -> &'static str {
172    if snap.session.is_some() {
173        return DEFAULT_FORMAT;
174    }
175    if snap.weekly.is_some() {
176        return WEEKLY_ONLY_FORMAT;
177    }
178    NO_WINDOWS_FORMAT
179}
180
181fn render_tooltip(
182    outcome: &VendorOutcome,
183    snap: &OpenAiSnapshot,
184    theme: &Theme,
185    now: DateTime<Utc>,
186) -> String {
187    let blue = &theme.blue;
188    let dim = &theme.dim;
189    let fg = &theme.fg;
190
191    let mut lines: Vec<TooltipLine> = Vec::new();
192    lines.push(TooltipLine::Center(format!(
193        "<span font_weight='bold' foreground='{blue}'>{plan}</span>",
194        plan = escape(&snap.plan)
195    )));
196    lines.push(TooltipLine::Sep);
197    lines.push(TooltipLine::Body("".into()));
198
199    if let Some(session) = snap.session.as_ref() {
200        push_window(&mut lines, "  󰔟  Codex 5h", session, theme, now, None);
201    }
202    if let Some(weekly) = snap.weekly.as_ref() {
203        if snap.session.is_some() {
204            lines.push(TooltipLine::Body("".into()));
205        }
206        push_window(&mut lines, "  󰃰  Codex weekly", weekly, theme, now, None);
207    }
208    if snap.session.is_none() && snap.weekly.is_none() {
209        lines.push(TooltipLine::Body(format!(
210            " <span foreground='{dim}'>no usage windows reported</span>"
211        )));
212    }
213
214    if let Some(cr) = snap.code_review.as_ref() {
215        lines.push(TooltipLine::Body("".into()));
216        push_window(
217            &mut lines,
218            "  󱦰  Code review (weekly)",
219            cr,
220            theme,
221            now,
222            None,
223        );
224    }
225
226    if let Some(c) = snap.credits.as_ref() {
227        lines.push(TooltipLine::Body("".into()));
228        lines.push(TooltipLine::Sep);
229        let label = if c.unlimited {
230            "unlimited".to_string()
231        } else {
232            c.balance.clone()
233        };
234        lines.push(TooltipLine::Body(format!(
235            " <span foreground='{fg}'>  󰄑  Credits</span>"
236        )));
237        lines.push(TooltipLine::Body(format!(
238            " <span foreground='{dim}'>     balance: {b}</span>",
239            b = escape(&label)
240        )));
241        if let Some((lo, hi)) = c.approx_local_messages {
242            lines.push(TooltipLine::Body(format!(
243                " <span foreground='{dim}'>     ~ {lo}-{hi} local messages</span>"
244            )));
245        }
246        if let Some((lo, hi)) = c.approx_cloud_messages {
247            lines.push(TooltipLine::Body(format!(
248                " <span foreground='{dim}'>     ~ {lo}-{hi} cloud messages</span>"
249            )));
250        }
251    }
252
253    if matches!(snap.source, OpenAiSource::Unavailable) {
254        lines.push(TooltipLine::Body("".into()));
255        lines.push(TooltipLine::Sep);
256        lines.push(TooltipLine::Body(format!(
257            " <span foreground='{dim}'>OpenAI plan usage requires Codex OAuth.</span>"
258        )));
259        lines.push(TooltipLine::Body(format!(
260            " <span foreground='{dim}'>Run `codex login` to enable.</span>"
261        )));
262    }
263
264    if let Some((code, msg)) = outcome.last_error.as_ref()
265        && *code != 0
266    {
267        let (icon, ecolor) = if *code >= 500 {
268            ("󰅚", theme.red.as_str())
269        } else {
270            ("󰀪", theme.orange.as_str())
271        };
272        lines.push(TooltipLine::Body("".into()));
273        lines.push(TooltipLine::Sep);
274        lines.push(TooltipLine::Body(format!(
275            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
276        )));
277        lines.push(TooltipLine::Body(format!(
278            "     <span foreground='{dim}'>{}</span>",
279            escape(msg)
280        )));
281    }
282
283    let updated = updated_at_hm(now, outcome.cache_age);
284    lines.push(TooltipLine::Body("".into()));
285    lines.push(TooltipLine::Sep);
286    lines.push(TooltipLine::Body(format!(
287        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
288    )));
289
290    render_bordered(&lines, theme)
291}
292
293impl From<FetchOutcome> for VendorOutcome {
294    fn from(o: FetchOutcome) -> Self {
295        Self {
296            snapshot: crate::usage::VendorSnapshot::Openai(o.snapshot),
297            stale: o.stale,
298            last_error: o.last_error,
299            cache_age: o.cache_age,
300        }
301    }
302}
303
304#[cfg(test)]
305mod tests {
306    use super::*;
307    use crate::usage::{OpenAiSnapshot, OpenAiSource, UsageWindow};
308
309    fn sample() -> OpenAiSnapshot {
310        OpenAiSnapshot {
311            plan: "ChatGPT Plus".into(),
312            session: Some(UsageWindow {
313                utilization_pct: 1,
314                resets_at: Some(Utc::now() + chrono::Duration::hours(5)),
315                window_duration: chrono::Duration::hours(5),
316            }),
317            weekly: Some(UsageWindow {
318                utilization_pct: 0,
319                resets_at: Some(Utc::now() + chrono::Duration::days(7)),
320                window_duration: chrono::Duration::days(7),
321            }),
322            code_review: None,
323            credits: None,
324            source: OpenAiSource::CodexOauth,
325        }
326    }
327
328    fn oc(s: OpenAiSnapshot) -> VendorOutcome {
329        VendorOutcome {
330            snapshot: crate::usage::VendorSnapshot::Openai(s),
331            stale: false,
332            last_error: None,
333            cache_age: Some(std::time::Duration::from_secs(15)),
334        }
335    }
336
337    fn opts() -> RenderOpts {
338        RenderOpts {
339            format: None,
340            tooltip_format: None,
341            icon: None,
342            pace_tolerance: 5,
343            format_pace_color: false,
344            tooltip_pace_pts: false,
345        }
346    }
347
348    #[test]
349    fn default_format_renders_session() {
350        let s = sample();
351        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
352        assert!(out.text.contains("1%"));
353    }
354
355    #[test]
356    fn dual_window_placeholders_keep_existing_values() {
357        let s = sample();
358        let values = build_placeholders(&s, &opts(), Utc::now());
359        assert_eq!(values["oai_session_pct"], "1");
360        assert_eq!(values["oai_weekly_pct"], "0");
361        assert_eq!(values["session_pct"], "1");
362        assert_eq!(values["weekly_pct"], "0");
363        // Cross-vendor elapsed aliases must mirror the oai_* values so the
364        // macOS menu bar can render pace markers for OpenAI.
365        assert_eq!(values["session_elapsed"], values["oai_session_elapsed"]);
366        assert_eq!(values["weekly_elapsed"], values["oai_weekly_elapsed"]);
367        assert!(!values["session_elapsed"].is_empty());
368        assert!(!values["weekly_elapsed"].is_empty());
369    }
370
371    #[test]
372    fn tooltip_has_both_windows() {
373        let s = sample();
374        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
375        assert!(out.tooltip.contains("Codex 5h"));
376        assert!(out.tooltip.contains("Codex weekly"));
377        assert!(!out.tooltip.contains("Code review"));
378        assert!(!out.tooltip.contains("Credits"));
379    }
380
381    #[test]
382    fn weekly_only_snapshot_uses_weekly_default_and_hides_session() {
383        let mut s = sample();
384        s.session = None;
385        s.weekly.as_mut().unwrap().utilization_pct = 66;
386        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
387        assert!(out.text.contains("66%"));
388        assert!(out.tooltip.contains("Codex weekly"));
389        assert!(!out.tooltip.contains("Codex 5h"));
390
391        let values = build_placeholders(&s, &opts(), Utc::now());
392        assert_eq!(values["oai_session_pct"], "");
393        assert_eq!(values["oai_weekly_pct"], "66");
394    }
395
396    #[test]
397    fn no_windows_uses_plan_only_format_and_reports_absence() {
398        let mut s = sample();
399        s.session = None;
400        s.weekly = None;
401        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
402        assert!(out.text.contains("ChatGPT Plus"));
403        assert!(!out.text.contains('%'));
404        assert!(out.tooltip.contains("no usage windows reported"));
405        assert!(!out.tooltip.contains("Codex 5h"));
406        assert!(!out.tooltip.contains("Codex weekly"));
407        // No windows means max utilization is 0, i.e. the lowest severity.
408        assert_eq!(severity(&s), PaceSeverity::Low);
409    }
410
411    #[test]
412    fn session_only_snapshot_renders_default_with_weekly_placeholders_empty() {
413        let mut s = sample();
414        s.weekly = None;
415        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
416        assert!(out.text.contains("1%"));
417        assert!(out.tooltip.contains("Codex 5h"));
418        assert!(!out.tooltip.contains("Codex weekly"));
419
420        let values = build_placeholders(&s, &opts(), Utc::now());
421        assert_eq!(values["oai_session_pct"], "1");
422        assert_eq!(values["oai_weekly_pct"], "");
423        assert_eq!(values["oai_weekly_reset"], "");
424        assert_eq!(values["oai_weekly_elapsed"], "");
425        assert_eq!(values["weekly_pct"], "");
426    }
427
428    #[test]
429    fn tooltip_includes_credits_block_when_present() {
430        let mut s = sample();
431        s.credits = Some(crate::usage::OpenAiCredits {
432            balance: "$5.00".into(),
433            has_credits: true,
434            unlimited: false,
435            approx_local_messages: Some((100, 200)),
436            approx_cloud_messages: Some((30, 50)),
437        });
438        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
439        assert!(out.tooltip.contains("Credits"));
440        assert!(out.tooltip.contains("$5.00"));
441        assert!(out.tooltip.contains("100-200 local messages"));
442        assert!(out.tooltip.contains("30-50 cloud messages"));
443    }
444
445    #[test]
446    fn unavailable_source_shows_codex_login_hint() {
447        let mut s = sample();
448        s.source = OpenAiSource::Unavailable;
449        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
450        assert!(out.tooltip.contains("codex login"));
451    }
452
453    #[test]
454    fn severity_picks_worst_window() {
455        let mut s = sample();
456        s.weekly.as_mut().unwrap().utilization_pct = 95;
457        assert_eq!(severity(&s), PaceSeverity::Critical);
458    }
459}