Skip to main content

ai_usagebar/opencode_go/
vendor.rs

1use std::collections::HashMap;
2use std::time::Duration;
3
4use chrono::{DateTime, Utc};
5
6use crate::countdown;
7use crate::format::{placeholders, substitute, updated_at_hm};
8use crate::pacing::PaceSeverity;
9use crate::pango::{color_span, escape, severity_color, severity_for};
10use crate::theme::Theme;
11use crate::tooltip::{Line as TooltipLine, render_bordered};
12use crate::vendor::{RenderOpts, VendorOutcome};
13use crate::waybar::{Class, WaybarOutput};
14
15use super::fetch::FetchOutcome;
16use super::types::{Usage, Window};
17
18pub const DEFAULT_FORMAT: &str = "{ocg_rolling_pct}% · {ocg_rolling_reset}";
19const DEFAULT_PLAN: &str = "OpenCode Go";
20const UNAVAILABLE: &str = "—";
21
22impl From<FetchOutcome> for VendorOutcome {
23    fn from(outcome: FetchOutcome) -> Self {
24        Self {
25            snapshot: crate::usage::VendorSnapshot::OpenCodeGo(outcome.snapshot),
26            stale: outcome.stale,
27            last_error: outcome.last_error,
28            cache_age: outcome.cache_age,
29        }
30    }
31}
32
33pub fn build_placeholders(usage: &Usage, now: DateTime<Utc>) -> HashMap<&'static str, String> {
34    build_placeholders_with_plan(DEFAULT_PLAN, usage, now)
35}
36
37pub fn build_placeholders_with_plan(
38    plan: &str,
39    usage: &Usage,
40    now: DateTime<Utc>,
41) -> HashMap<&'static str, String> {
42    let plan = sanitize(plan);
43    let rolling = window_values(usage.rolling.as_ref(), now);
44    let weekly = window_values(usage.weekly.as_ref(), now);
45    let monthly = window_values(usage.monthly.as_ref(), now);
46
47    placeholders([
48        ("vendor_short", "ocg".to_string()),
49        ("plan", plan.clone()),
50        ("ocg_plan", plan),
51        ("session_pct", rolling.percent.clone()),
52        ("session_reset", rolling.reset.clone()),
53        ("weekly_pct", weekly.percent.clone()),
54        ("weekly_reset", weekly.reset.clone()),
55        ("ocg_rolling_pct", rolling.percent),
56        ("ocg_rolling_reset", rolling.reset),
57        ("ocg_rolling_status", rolling.status),
58        ("ocg_weekly_pct", weekly.percent),
59        ("ocg_weekly_reset", weekly.reset),
60        ("ocg_weekly_status", weekly.status),
61        ("ocg_monthly_pct", monthly.percent),
62        ("ocg_monthly_reset", monthly.reset),
63        ("ocg_monthly_status", monthly.status),
64    ])
65}
66
67#[derive(Debug)]
68struct WindowValues {
69    percent: String,
70    reset: String,
71    status: String,
72}
73
74fn window_values(window: Option<&Window>, now: DateTime<Utc>) -> WindowValues {
75    let Some(window) = window else {
76        return WindowValues {
77            percent: UNAVAILABLE.to_string(),
78            reset: UNAVAILABLE.to_string(),
79            status: UNAVAILABLE.to_string(),
80        };
81    };
82    WindowValues {
83        percent: window.percent.to_string(),
84        reset: countdown::format(Some(window.resets_at), now),
85        status: sanitize(&window.status),
86    }
87}
88
89fn sanitize(value: &str) -> String {
90    crate::display::sanitize_untrusted_field(value)
91}
92
93pub fn severity(usage: &Usage) -> PaceSeverity {
94    usage
95        .rolling
96        .iter()
97        .chain(usage.weekly.iter())
98        .chain(usage.monthly.iter())
99        .map(|window| window.percent as i32)
100        .max()
101        .map(severity_for)
102        .unwrap_or(PaceSeverity::Low)
103}
104
105/// Renderer shape matches the existing vendor adapters. `snap` remains local
106/// because the shared enum does not yet have an OpenCode-Go arm.
107pub fn render(
108    outcome: &VendorOutcome,
109    snap: &Usage,
110    theme: &Theme,
111    opts: &RenderOpts,
112    now: DateTime<Utc>,
113) -> WaybarOutput {
114    render_with_meta(
115        snap,
116        outcome.stale,
117        outcome.last_error.as_ref(),
118        outcome.cache_age,
119        theme,
120        opts,
121        now,
122    )
123}
124
125fn render_with_meta(
126    snap: &Usage,
127    stale: bool,
128    last_error: Option<&(u16, String)>,
129    cache_age: Option<Duration>,
130    theme: &Theme,
131    opts: &RenderOpts,
132    now: DateTime<Utc>,
133) -> WaybarOutput {
134    let sev = severity(snap);
135    let format = opts.format.as_deref().unwrap_or(DEFAULT_FORMAT);
136    let values = escaped_placeholders(snap, now);
137    let mut text = substitute(format, &values);
138    if stale {
139        text.push_str(" ⏸");
140    }
141    let icon_prefix = match opts.icon.as_deref() {
142        Some(icon) if !icon.is_empty() => format!("{} ", escape(icon)),
143        _ => String::new(),
144    };
145    let bar_text = color_span(severity_color(sev, theme), &format!("{icon_prefix}{text}"));
146    let tooltip = opts
147        .tooltip_format
148        .as_deref()
149        .map(|format| substitute(format, &values))
150        .unwrap_or_else(|| render_tooltip(snap, stale, last_error, cache_age, theme, now));
151
152    WaybarOutput {
153        text: bar_text,
154        tooltip,
155        class: Class::from(sev),
156    }
157}
158
159fn escaped_placeholders(usage: &Usage, now: DateTime<Utc>) -> HashMap<&'static str, String> {
160    let mut values = build_placeholders(usage, now);
161    for key in [
162        "plan",
163        "ocg_plan",
164        "ocg_rolling_status",
165        "ocg_weekly_status",
166        "ocg_monthly_status",
167    ] {
168        if let Some(value) = values.get_mut(key) {
169            *value = escape(value);
170        }
171    }
172    values
173}
174
175fn render_tooltip(
176    snap: &Usage,
177    stale: bool,
178    last_error: Option<&(u16, String)>,
179    cache_age: Option<Duration>,
180    theme: &Theme,
181    now: DateTime<Utc>,
182) -> String {
183    let mut lines = vec![TooltipLine::Center(format!(
184        "<span font_weight='bold' foreground='{}'>{}</span>",
185        theme.blue,
186        escape(DEFAULT_PLAN)
187    ))];
188    lines.push(TooltipLine::Sep);
189    lines.push(TooltipLine::Body(String::new()));
190
191    let mut present = false;
192    for (label, window) in [
193        ("Rolling", snap.rolling.as_ref()),
194        ("Weekly", snap.weekly.as_ref()),
195        ("Monthly", snap.monthly.as_ref()),
196    ] {
197        let Some(window) = window else {
198            continue;
199        };
200        present = true;
201        let values = window_values(Some(window), now);
202        lines.push(TooltipLine::Body(format!(
203            "  {}  {}% · {} · {}",
204            label,
205            escape(&values.percent),
206            escape(&values.reset),
207            escape(&values.status)
208        )));
209    }
210    if !present {
211        lines.push(TooltipLine::Body(format!(
212            " <span foreground='{}'>no usage windows reported</span>",
213            theme.dim
214        )));
215    }
216    if stale {
217        lines.push(TooltipLine::Body(String::new()));
218        lines.push(TooltipLine::Body(format!(
219            " <span foreground='{}'>  ⏸  Showing cached data</span>",
220            theme.orange
221        )));
222    }
223    if let Some((code, message)) = last_error
224        && *code != 0
225    {
226        lines.push(TooltipLine::Body(String::new()));
227        lines.push(TooltipLine::Sep);
228        lines.push(TooltipLine::Body(format!(
229            " <span foreground='{}'>  HTTP {code}: {}</span>",
230            theme.orange,
231            escape(message)
232        )));
233    }
234
235    lines.push(TooltipLine::Body(String::new()));
236    lines.push(TooltipLine::Sep);
237    lines.push(TooltipLine::Body(format!(
238        " <span foreground='{}'>  Updated {}</span>",
239        theme.dim,
240        updated_at_hm(now, cache_age)
241    )));
242    render_bordered(&lines, theme)
243}
244
245#[cfg(test)]
246mod tests {
247    use chrono::{DateTime, Utc};
248
249    use super::*;
250    use crate::opencode_go::types::{Usage, Window};
251
252    fn at(value: &str) -> DateTime<Utc> {
253        value.parse().expect("RFC3339 timestamp")
254    }
255
256    fn sample_usage() -> Usage {
257        Usage {
258            rolling: Some(Window {
259                status: "ok".into(),
260                percent: 12.3,
261                resets_at: at("2026-08-16T20:00:00Z"),
262            }),
263            weekly: Some(Window {
264                status: "rate-limited".into(),
265                percent: 45.6,
266                resets_at: at("2026-08-20T00:00:00Z"),
267            }),
268            monthly: Some(Window {
269                status: "ok".into(),
270                percent: 78.9,
271                resets_at: at("2026-09-01T00:00:00Z"),
272            }),
273        }
274    }
275
276    #[test]
277    fn exposes_exact_opencode_go_and_generic_placeholders() {
278        let values = build_placeholders(&sample_usage(), at("2026-08-16T18:00:00Z"));
279
280        assert_eq!(values["vendor_short"], "ocg");
281        assert_eq!(values["session_pct"], "12.3");
282        assert_eq!(values["weekly_pct"], "45.6");
283        assert_eq!(values["ocg_rolling_pct"], "12.3");
284        assert_eq!(values["ocg_rolling_status"], "ok");
285        assert_eq!(values["ocg_weekly_pct"], "45.6");
286        assert_eq!(values["ocg_weekly_status"], "rate-limited");
287        assert_eq!(values["ocg_monthly_pct"], "78.9");
288        assert_eq!(values["ocg_monthly_status"], "ok");
289    }
290
291    #[test]
292    fn default_format_is_rolling_percentage_and_reset() {
293        assert_eq!(DEFAULT_FORMAT, "{ocg_rolling_pct}% · {ocg_rolling_reset}");
294    }
295
296    #[test]
297    fn absent_windows_are_unavailable_not_zero() {
298        let values = build_placeholders(
299            &Usage {
300                rolling: None,
301                weekly: None,
302                monthly: None,
303            },
304            at("2026-08-16T18:00:00Z"),
305        );
306
307        for key in [
308            "session_pct",
309            "weekly_pct",
310            "ocg_rolling_pct",
311            "ocg_weekly_pct",
312            "ocg_monthly_pct",
313        ] {
314            assert_eq!(values[key], "—", "{key} should be unavailable");
315            assert_ne!(values[key], "0");
316        }
317        for key in [
318            "session_reset",
319            "weekly_reset",
320            "ocg_rolling_reset",
321            "ocg_weekly_reset",
322            "ocg_monthly_reset",
323            "ocg_rolling_status",
324            "ocg_weekly_status",
325            "ocg_monthly_status",
326        ] {
327            assert_eq!(values[key], "—", "{key} should be unavailable");
328        }
329    }
330
331    #[test]
332    fn plan_and_status_are_sanitized() {
333        let usage = Usage {
334            rolling: Some(Window {
335                status: "ok\u{1b}[31m\u{7}".into(),
336                percent: 1.0,
337                resets_at: at("2026-08-16T20:00:00Z"),
338            }),
339            weekly: None,
340            monthly: None,
341        };
342        let values = build_placeholders_with_plan(
343            "OpenCode\u{1b}[31m Go",
344            &usage,
345            at("2026-08-16T18:00:00Z"),
346        );
347
348        assert!(!values["plan"].contains('\u{1b}'));
349        assert!(!values["ocg_rolling_status"].contains('\u{1b}'));
350        assert!(!values["ocg_rolling_status"].contains('\u{7}'));
351    }
352}