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