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::{self, color_span, escape, severity_color, severity_for};
12use crate::theme::Theme;
13use crate::tooltip::{Line as TooltipLine, 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);
179    lines.push(TooltipLine::Body("".into()));
180    push_window(&mut lines, "  󰃰  Codex weekly", &snap.weekly, theme, now);
181
182    if let Some(cr) = snap.code_review.as_ref() {
183        lines.push(TooltipLine::Body("".into()));
184        push_window(&mut lines, "  󱦰  Code review (weekly)", cr, theme, now);
185    }
186
187    if let Some(c) = snap.credits.as_ref() {
188        lines.push(TooltipLine::Body("".into()));
189        lines.push(TooltipLine::Sep);
190        let label = if c.unlimited {
191            "unlimited".to_string()
192        } else {
193            c.balance.clone()
194        };
195        lines.push(TooltipLine::Body(format!(
196            " <span foreground='{fg}'>  󰄑  Credits</span>"
197        )));
198        lines.push(TooltipLine::Body(format!(
199            " <span foreground='{dim}'>     balance: {b}</span>",
200            b = escape(&label)
201        )));
202        if let Some((lo, hi)) = c.approx_local_messages {
203            lines.push(TooltipLine::Body(format!(
204                " <span foreground='{dim}'>     ~ {lo}-{hi} local messages</span>"
205            )));
206        }
207        if let Some((lo, hi)) = c.approx_cloud_messages {
208            lines.push(TooltipLine::Body(format!(
209                " <span foreground='{dim}'>     ~ {lo}-{hi} cloud messages</span>"
210            )));
211        }
212    }
213
214    if matches!(snap.source, OpenAiSource::Unavailable) {
215        lines.push(TooltipLine::Body("".into()));
216        lines.push(TooltipLine::Sep);
217        lines.push(TooltipLine::Body(format!(
218            " <span foreground='{dim}'>OpenAI plan usage requires Codex OAuth.</span>"
219        )));
220        lines.push(TooltipLine::Body(format!(
221            " <span foreground='{dim}'>Run `codex login` to enable.</span>"
222        )));
223    }
224
225    if let Some((code, msg)) = outcome.last_error.as_ref()
226        && *code != 0
227    {
228        let (icon, ecolor) = if *code >= 500 {
229            ("󰅚", theme.red.as_str())
230        } else {
231            ("󰀪", theme.orange.as_str())
232        };
233        lines.push(TooltipLine::Body("".into()));
234        lines.push(TooltipLine::Sep);
235        lines.push(TooltipLine::Body(format!(
236            " <span foreground='{ecolor}'>  {icon}  HTTP {code}</span>"
237        )));
238        lines.push(TooltipLine::Body(format!(
239            "     <span foreground='{dim}'>{}</span>",
240            escape(msg)
241        )));
242    }
243
244    let updated = updated_at_hm(now, outcome.cache_age);
245    lines.push(TooltipLine::Body("".into()));
246    lines.push(TooltipLine::Sep);
247    lines.push(TooltipLine::Body(format!(
248        " <span foreground='{dim}'>  󰅐  Updated {updated}</span>"
249    )));
250
251    render_bordered(&lines, theme)
252}
253
254fn push_window(
255    lines: &mut Vec<TooltipLine>,
256    label: &str,
257    w: &crate::usage::UsageWindow,
258    theme: &Theme,
259    now: DateTime<Utc>,
260) {
261    let color = severity_color(severity_for(w.utilization_pct), theme);
262    let bar = pango::progress_bar(w.utilization_pct, color, theme, None);
263    let fg = &theme.fg;
264    let dim = &theme.dim;
265    lines.push(TooltipLine::Body(format!(
266        " <span foreground='{fg}'>{label}</span>"
267    )));
268    lines.push(TooltipLine::Body(format!(
269        "   {bar}  <span font_weight='bold' foreground='{color}'>{pct}%</span>",
270        pct = w.utilization_pct
271    )));
272    lines.push(TooltipLine::Body(format!(
273        " <span foreground='{dim}'>  ⏱  Resets in {cd}</span>",
274        cd = escape(&countdown::format(w.resets_at, now))
275    )));
276}
277
278impl From<FetchOutcome> for VendorOutcome {
279    fn from(o: FetchOutcome) -> Self {
280        Self {
281            snapshot: crate::usage::VendorSnapshot::Openai(o.snapshot),
282            stale: o.stale,
283            last_error: o.last_error,
284            cache_age: o.cache_age,
285        }
286    }
287}
288
289#[cfg(test)]
290mod tests {
291    use super::*;
292    use crate::usage::{OpenAiSnapshot, OpenAiSource, UsageWindow};
293
294    fn sample() -> OpenAiSnapshot {
295        OpenAiSnapshot {
296            plan: "ChatGPT Plus".into(),
297            session: UsageWindow {
298                utilization_pct: 1,
299                resets_at: Some(Utc::now() + chrono::Duration::hours(5)),
300                window_duration: chrono::Duration::hours(5),
301            },
302            weekly: UsageWindow {
303                utilization_pct: 0,
304                resets_at: Some(Utc::now() + chrono::Duration::days(7)),
305                window_duration: chrono::Duration::days(7),
306            },
307            code_review: None,
308            credits: None,
309            source: OpenAiSource::CodexOauth,
310        }
311    }
312
313    fn oc(s: OpenAiSnapshot) -> VendorOutcome {
314        VendorOutcome {
315            snapshot: crate::usage::VendorSnapshot::Openai(s),
316            stale: false,
317            last_error: None,
318            cache_age: Some(std::time::Duration::from_secs(15)),
319        }
320    }
321
322    fn opts() -> RenderOpts {
323        RenderOpts {
324            format: None,
325            tooltip_format: None,
326            icon: None,
327            pace_tolerance: 5,
328            format_pace_color: false,
329            tooltip_pace_pts: false,
330        }
331    }
332
333    #[test]
334    fn default_format_renders_session() {
335        let s = sample();
336        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
337        assert!(out.text.contains("1%"));
338    }
339
340    #[test]
341    fn tooltip_has_both_windows() {
342        let s = sample();
343        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
344        assert!(out.tooltip.contains("Codex 5h"));
345        assert!(out.tooltip.contains("Codex weekly"));
346        assert!(!out.tooltip.contains("Code review"));
347        assert!(!out.tooltip.contains("Credits"));
348    }
349
350    #[test]
351    fn tooltip_includes_credits_block_when_present() {
352        let mut s = sample();
353        s.credits = Some(crate::usage::OpenAiCredits {
354            balance: "$5.00".into(),
355            has_credits: true,
356            unlimited: false,
357            approx_local_messages: Some((100, 200)),
358            approx_cloud_messages: Some((30, 50)),
359        });
360        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
361        assert!(out.tooltip.contains("Credits"));
362        assert!(out.tooltip.contains("$5.00"));
363        assert!(out.tooltip.contains("100-200 local messages"));
364        assert!(out.tooltip.contains("30-50 cloud messages"));
365    }
366
367    #[test]
368    fn unavailable_source_shows_codex_login_hint() {
369        let mut s = sample();
370        s.source = OpenAiSource::Unavailable;
371        let out = render(&oc(s.clone()), &s, &Theme::default(), &opts(), Utc::now());
372        assert!(out.tooltip.contains("codex login"));
373    }
374
375    #[test]
376    fn severity_picks_worst_window() {
377        let mut s = sample();
378        s.weekly.utilization_pct = 95;
379        assert_eq!(severity(&s), PaceSeverity::Critical);
380    }
381}