Skip to main content

ai_usagebar/copilot/
vendor.rs

1//! GitHub Copilot Waybar renderer.
2
3use std::collections::HashMap;
4
5use chrono::{DateTime, Utc};
6
7use crate::countdown;
8use crate::format::{placeholders, substitute, updated_at_hm};
9use crate::pacing::PaceSeverity;
10use crate::pango::{color_span, escape, severity_color, severity_for};
11use crate::theme::Theme;
12use crate::tooltip::{Line as TooltipLine, render_bordered};
13use crate::vendor::{RenderOpts, VendorId, VendorOutcome};
14use crate::waybar::{Class, WaybarOutput};
15
16use super::fetch::FetchOutcome;
17use super::types::Snapshot;
18
19pub const DEFAULT_FORMAT: &str = "{copilot_premium_pct}% · {copilot_reset}";
20const UNAVAILABLE: &str = "—";
21
22impl From<FetchOutcome> for VendorOutcome {
23    fn from(outcome: FetchOutcome) -> Self {
24        outcome.map(crate::usage::VendorSnapshot::Copilot)
25    }
26}
27
28pub fn build_placeholders(snap: &Snapshot, now: DateTime<Utc>) -> HashMap<&'static str, String> {
29    let premium = quota_values(snap.premium.as_ref());
30    let chat = quota_values(snap.chat.as_ref());
31    let completions = quota_values(snap.completions.as_ref());
32    let reset = countdown::format(snap.reset_at, now);
33    placeholders([
34        ("icon", "󰊤".to_string()),
35        ("vendor_short", VendorId::Copilot.short_name().to_string()),
36        ("plan", crate::display::sanitize_untrusted_field(&snap.plan)),
37        ("session_pct", premium.percent.clone()),
38        ("session_reset", reset.clone()),
39        ("weekly_pct", chat.percent.clone()),
40        ("weekly_reset", reset.clone()),
41        (
42            "copilot_plan",
43            crate::display::sanitize_untrusted_field(&snap.plan),
44        ),
45        ("copilot_reset", reset),
46        ("copilot_premium_pct", premium.percent),
47        ("copilot_premium_used", premium.used),
48        ("copilot_premium_limit", premium.limit),
49        ("copilot_chat_pct", chat.percent),
50        ("copilot_chat_used", chat.used),
51        ("copilot_chat_limit", chat.limit),
52        ("copilot_completions_pct", completions.percent),
53        ("copilot_completions_used", completions.used),
54        ("copilot_completions_limit", completions.limit),
55    ])
56}
57
58struct QuotaValues {
59    percent: String,
60    used: String,
61    limit: String,
62}
63
64fn quota_values(quota: Option<&super::types::Quota>) -> QuotaValues {
65    let Some(quota) = quota else {
66        return QuotaValues {
67            percent: UNAVAILABLE.into(),
68            used: UNAVAILABLE.into(),
69            limit: UNAVAILABLE.into(),
70        };
71    };
72    if quota.unlimited {
73        return QuotaValues {
74            percent: "0".into(),
75            used: "0".into(),
76            limit: "unlimited".into(),
77        };
78    }
79    let (used, limit) = quota
80        .used_and_entitlement()
81        .map(|(used, limit)| (used.to_string(), limit.to_string()))
82        .unwrap_or_else(|| (UNAVAILABLE.into(), UNAVAILABLE.into()));
83    QuotaValues {
84        percent: quota.used_pct().to_string(),
85        used,
86        limit,
87    }
88}
89
90pub fn severity(snap: &Snapshot) -> PaceSeverity {
91    severity_for(snap.worst_pct())
92}
93
94pub fn render(
95    outcome: &VendorOutcome,
96    snap: &Snapshot,
97    theme: &Theme,
98    opts: &RenderOpts,
99    now: DateTime<Utc>,
100) -> WaybarOutput {
101    let severity = severity(snap);
102    let format = opts.format.as_deref().unwrap_or(DEFAULT_FORMAT);
103    let mut values = build_placeholders(snap, now);
104    for key in ["plan", "copilot_plan"] {
105        if let Some(value) = values.get_mut(key) {
106            *value = escape(value);
107        }
108    }
109    let mut text = substitute(format, &values);
110    if outcome.stale {
111        text.push_str(" ⏸");
112    }
113    let icon = opts
114        .icon
115        .as_deref()
116        .filter(|icon| !icon.is_empty())
117        .map(|icon| format!("{} ", escape(icon)))
118        .unwrap_or_default();
119    let tooltip = opts
120        .tooltip_format
121        .as_deref()
122        .map(|format| substitute(format, &values))
123        .unwrap_or_else(|| render_tooltip(outcome, snap, theme, now));
124    WaybarOutput {
125        text: color_span(severity_color(severity, theme), &format!("{icon}{text}")),
126        tooltip,
127        class: Class::from(severity),
128    }
129}
130
131fn render_tooltip(
132    outcome: &VendorOutcome,
133    snap: &Snapshot,
134    theme: &Theme,
135    now: DateTime<Utc>,
136) -> String {
137    let mut lines = vec![TooltipLine::Center(format!(
138        "<span font_weight='bold' foreground='{}'>GitHub Copilot {}</span>",
139        theme.blue,
140        escape(&crate::display::sanitize_untrusted_field(&snap.plan))
141    ))];
142    lines.push(TooltipLine::Sep);
143    lines.push(TooltipLine::Body(String::new()));
144    for (label, quota) in snap.quotas() {
145        let usage = if quota.unlimited {
146            "Unlimited".to_string()
147        } else if let Some((used, entitlement)) = quota.used_and_entitlement() {
148            format!("{}% · {used} of {entitlement} used", quota.used_pct())
149        } else {
150            format!(
151                "{}% · {}% remaining",
152                quota.used_pct(),
153                quota.percent_remaining
154            )
155        };
156        lines.push(TooltipLine::Body(format!("  {label}  {}", escape(&usage))));
157    }
158    lines.push(TooltipLine::Body(format!(
159        "  Resets  {}",
160        escape(&countdown::format(snap.reset_at, now))
161    )));
162    if outcome.stale {
163        lines.push(TooltipLine::Body(String::new()));
164        lines.push(TooltipLine::Body(format!(
165            " <span foreground='{}'>  ⏸  Showing cached data</span>",
166            theme.orange
167        )));
168    }
169    if let Some((code, message)) = outcome.last_error.as_ref()
170        && *code != 0
171    {
172        lines.push(TooltipLine::Body(String::new()));
173        lines.push(TooltipLine::Sep);
174        lines.push(TooltipLine::Body(format!(
175            " <span foreground='{}'>  HTTP {code}: {}</span>",
176            theme.orange,
177            escape(message)
178        )));
179    }
180    lines.push(TooltipLine::Body(String::new()));
181    lines.push(TooltipLine::Sep);
182    lines.push(TooltipLine::Body(format!(
183        " <span foreground='{}'>  Updated {}</span>",
184        theme.dim,
185        updated_at_hm(now, outcome.cache_age)
186    )));
187    render_bordered(&lines, theme)
188}
189
190#[cfg(test)]
191mod tests {
192    use super::*;
193    use crate::copilot::types::Quota;
194
195    fn sample() -> Snapshot {
196        Snapshot {
197            plan: "Pro".into(),
198            premium: Some(Quota {
199                percent_remaining: 15,
200                entitlement: Some(300),
201                remaining: Some(45),
202                unlimited: false,
203            }),
204            chat: None,
205            completions: None,
206            reset_at: None,
207        }
208    }
209
210    #[test]
211    fn exposes_provider_and_generic_quota_placeholders() {
212        let values = build_placeholders(&sample(), Utc::now());
213        assert_eq!(values["vendor_short"], "ghc");
214        assert_eq!(values["copilot_premium_pct"], "85");
215        assert_eq!(values["copilot_premium_used"], "255");
216        assert_eq!(values["session_pct"], "85");
217        assert_eq!(values["weekly_pct"], UNAVAILABLE);
218    }
219
220    #[test]
221    fn renderer_uses_the_premium_quota_and_canonical_provider_name() {
222        let snap = sample();
223        let outcome = VendorOutcome::fresh(crate::usage::VendorSnapshot::Copilot(snap.clone()));
224        let output = render(
225            &outcome,
226            &snap,
227            &Theme::default(),
228            &RenderOpts {
229                format: None,
230                tooltip_format: None,
231                icon: None,
232                pace_tolerance: 5,
233                format_pace_color: false,
234                tooltip_pace_pts: false,
235            },
236            Utc::now(),
237        );
238        assert!(output.text.contains("85%"));
239        assert!(output.tooltip.contains("GitHub Copilot Pro"));
240    }
241}