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