Skip to main content

cmd_usage/arkuse/
api.rs

1// The ARK control plane, through `arkcli`.
2//
3// Unit of truth: AFP. `arkcli usage plan` answers the rolling windows
4// (5h / daily / weekly / monthly) with `used`, `total`, `percent` and the
5// server's own `reset_at` — the same numbers the Volcengine console shows, and
6// the only view that can see traffic from another machine or another harness.
7//
8// Two design rules, both learned on the OpenCode side of this crate:
9//
10// * the *percent* the server hands us is printed verbatim, never re-rounded;
11// * the local estimate never claims to be the meter. When `arkcli` cannot
12//   answer, `Unavailable::reason` says why and the renderers label the rows
13//   `estimated locally`.
14//
15// Shelling out, not HTTP: `arkcli` already owns the credential chain (SSO,
16// profile pinning, region/project selection), and re-implementing it here would
17// duplicate a signed control-plane client for no gain. One process per poll,
18// bounded by `CLI_TIMEOUT_SECS` so a hung CLI cannot pin a watch frame.
19use crate::arkuse::window::days_from_civil;
20use serde::Deserialize;
21use std::process::{Command, Stdio};
22use std::time::{Duration, Instant};
23
24/// Hard cap on one `arkcli` call. The control plane answers in ~2s; anything
25/// past this is a hang, and the frame must repaint with the local estimate.
26pub const CLI_TIMEOUT_SECS: u64 = 20;
27
28#[derive(Debug, Clone)]
29pub struct Meter {
30    /// `ok` / `exceeded` — ARK reports no status, so it is derived from percent.
31    pub status: String,
32    /// The server's own percentage, as served.
33    pub percent: f64,
34    pub used: f64,
35    pub quota: f64,
36    /// The raw instant, for JSON and the header.
37    pub resets_at: Option<String>,
38    reset_ms: Option<i64>,
39    offset_secs: i64,
40}
41
42impl Meter {
43    pub fn from_parts(
44        status: &str,
45        percent: f64,
46        used: f64,
47        quota: f64,
48        resets_at: Option<&str>,
49        offset_secs: Option<i64>,
50    ) -> Self {
51        let parsed = resets_at.and_then(parse_instant);
52        Self {
53            status: status.into(),
54            percent,
55            used,
56            quota,
57            resets_at: resets_at.map(str::to_string),
58            reset_ms: parsed.map(|(ms, _)| ms),
59            offset_secs: offset_secs
60                .or(parsed.map(|(_, off)| off))
61                .unwrap_or(8 * 3600),
62        }
63    }
64
65    pub fn reset_ms(&self) -> Option<i64> {
66        self.reset_ms
67    }
68
69    /// The server's UTC offset for this window, seconds east-positive.
70    pub fn offset_secs(&self) -> i64 {
71        self.offset_secs
72    }
73
74    pub fn exceeded(&self) -> bool {
75        self.percent >= 100.0
76    }
77
78    /// AFP still available in this window, floored at zero.
79    pub fn remaining(&self) -> f64 {
80        (self.quota - self.used).max(0.0)
81    }
82}
83
84#[derive(Debug, Clone)]
85pub struct Usage {
86    pub five_hour: Meter,
87    /// Present only for plans/workloads that expose a daily window (media and
88    /// Harness items); a text-only workload answers 5h/weekly/monthly.
89    pub daily: Option<Meter>,
90    pub weekly: Meter,
91    pub monthly: Meter,
92    /// The tier the subscription is on, from the control plane itself — the one
93    /// fact no local derivation can know.
94    pub tier: Option<String>,
95    pub profile: String,
96}
97
98#[derive(Debug, Clone, PartialEq, Eq)]
99pub enum Unavailable {
100    /// `--no-api`: deliberate, not a fault.
101    Disabled,
102    /// No `agent-plan` profile and no `--profile`.
103    NoProfile,
104    /// `arkcli` is not on PATH.
105    CliMissing,
106    Failed(String),
107}
108
109impl Unavailable {
110    pub fn reason(&self) -> String {
111        match self {
112            Unavailable::Disabled => "API disabled (--no-api)".into(),
113            Unavailable::NoProfile => "no agent-plan arkcli profile".into(),
114            Unavailable::CliMissing => "arkcli not found on PATH".into(),
115            Unavailable::Failed(why) => format!("arkcli: {why}"),
116        }
117    }
118}
119
120#[derive(Deserialize)]
121struct ProfileList {
122    profiles: Vec<ProfileEntry>,
123}
124
125#[derive(Deserialize)]
126struct ProfileEntry {
127    name: String,
128    #[serde(rename = "type")]
129    kind: Option<String>,
130    #[serde(rename = "plan_tier")]
131    tier: Option<String>,
132}
133
134/// Resolve the profile to query: `--profile`, then `$ARKUSE_PROFILE`, then the
135/// first `agent-plan` profile arkcli lists. Also returns that profile's tier,
136/// which is the default for `--plan`.
137pub fn resolve(explicit: Option<&str>) -> Option<(String, Option<String>)> {
138    let name = explicit.map(str::to_string).or_else(|| {
139        std::env::var("ARKUSE_PROFILE")
140            .ok()
141            .filter(|v| !v.is_empty())
142    });
143    let listed = profiles();
144    if let Some(name) = name {
145        let tier = listed
146            .iter()
147            .find(|(n, _, _)| n == &name)
148            .and_then(|(_, _, t)| t.clone());
149        return Some((name, tier));
150    }
151    listed
152        .into_iter()
153        .find(|(_, kind, _)| kind.as_deref() == Some("agent-plan"))
154        .map(|(name, _, tier)| (name, tier))
155}
156
157fn profiles() -> Vec<(String, Option<String>, Option<String>)> {
158    let Ok(out) = run_cli(&["profile", "list", "--format", "json"]) else {
159        return Vec::new();
160    };
161    serde_json::from_str::<ProfileList>(&out)
162        .map(|list| {
163            list.profiles
164                .into_iter()
165                .map(|p| (p.name, p.kind, p.tier))
166                .collect()
167        })
168        .unwrap_or_default()
169}
170
171#[derive(Deserialize)]
172struct UsageEnvelope {
173    items: Vec<UsageItem>,
174}
175
176#[derive(Deserialize)]
177struct UsageItem {
178    edition: Option<String>,
179    tier: Option<String>,
180    subscribed: Option<bool>,
181    #[serde(default)]
182    periods: Vec<Period>,
183}
184
185#[derive(Deserialize)]
186struct Period {
187    label: String,
188    used: f64,
189    total: f64,
190    percent: f64,
191    reset_at: Option<String>,
192}
193
194/// Ask the control plane for the plan's meters.
195pub fn fetch(profile: &str) -> Result<Usage, Unavailable> {
196    let raw = run_cli(&[
197        "usage",
198        "plan",
199        "--profile",
200        profile,
201        "--product",
202        "agent-plan",
203        "--format",
204        "json",
205    ])?;
206    parse_usage(&raw, profile)
207}
208
209fn parse_usage(raw: &str, profile: &str) -> Result<Usage, Unavailable> {
210    let envelope: UsageEnvelope =
211        serde_json::from_str(raw).map_err(|e| Unavailable::Failed(format!("parse: {e}")))?;
212    // Prefer the personal edition when the account holds several SKUs.
213    let item = envelope
214        .items
215        .iter()
216        .find(|i| i.edition.as_deref() == Some("personal"))
217        .or_else(|| envelope.items.first())
218        .ok_or_else(|| Unavailable::Failed("no plan items in the answer".into()))?;
219    if item.subscribed == Some(false) {
220        return Err(Unavailable::Failed("no active subscription".into()));
221    }
222    let meter = |label: &str| -> Option<Meter> {
223        item.periods.iter().find(|p| p.label == label).map(|p| {
224            let status = if p.percent >= 100.0 { "exceeded" } else { "ok" };
225            Meter::from_parts(
226                status,
227                p.percent,
228                p.used,
229                p.total,
230                p.reset_at.as_deref(),
231                None,
232            )
233        })
234    };
235    let missing = |label: &str| Unavailable::Failed(format!("the answer has no {label} window"));
236    Ok(Usage {
237        five_hour: meter("5h").ok_or_else(|| missing("5h"))?,
238        daily: meter("daily"),
239        weekly: meter("weekly").ok_or_else(|| missing("weekly"))?,
240        monthly: meter("monthly").ok_or_else(|| missing("monthly"))?,
241        tier: item.tier.clone(),
242        profile: profile.into(),
243    })
244}
245
246/// Run `arkcli` with a hard timeout, returning stdout. Unix-only process kill is
247/// fine here: this crate already requires a Unix terminal for its watch mode.
248fn run_cli(args: &[&str]) -> Result<String, Unavailable> {
249    let mut child = Command::new("arkcli")
250        .args(args)
251        .stdin(Stdio::null())
252        .stdout(Stdio::piped())
253        .stderr(Stdio::piped())
254        .spawn()
255        .map_err(|e| {
256            if e.kind() == std::io::ErrorKind::NotFound {
257                Unavailable::CliMissing
258            } else {
259                Unavailable::Failed(e.to_string())
260            }
261        })?;
262    let started = Instant::now();
263    loop {
264        match child.try_wait() {
265            Ok(Some(status)) => {
266                let out = child
267                    .wait_with_output()
268                    .map_err(|e| Unavailable::Failed(format!("read stdout: {e}")))?;
269                if !status.success() {
270                    let err = String::from_utf8_lossy(&out.stderr);
271                    return Err(Unavailable::Failed(cli_error(&err)));
272                }
273                return Ok(String::from_utf8_lossy(&out.stdout).into_owned());
274            }
275            Ok(None) => {
276                if started.elapsed() > Duration::from_secs(CLI_TIMEOUT_SECS) {
277                    let _ = child.kill();
278                    return Err(Unavailable::Failed(format!(
279                        "timed out after {CLI_TIMEOUT_SECS}s"
280                    )));
281                }
282                std::thread::sleep(Duration::from_millis(50));
283            }
284            Err(e) => return Err(Unavailable::Failed(e.to_string())),
285        }
286    }
287}
288
289/// `arkcli` reports failures as a JSON envelope on stderr
290/// (`{"ok":false,"error":{"message":…}}`), so the first line is just `{`. Take
291/// the message when there is one, else the last non-empty line — the actionable
292/// sentence, not the brace.
293fn cli_error(stderr: &str) -> String {
294    #[derive(Deserialize)]
295    struct Envelope {
296        error: Option<Inner>,
297    }
298    #[derive(Deserialize)]
299    struct Inner {
300        message: Option<String>,
301    }
302    if let Ok(env) = serde_json::from_str::<Envelope>(stderr.trim()) {
303        if let Some(text) = env.error.and_then(|e| e.message) {
304            if !text.trim().is_empty() {
305                return text.trim().to_string();
306            }
307        }
308    }
309    stderr
310        .lines()
311        .map(str::trim)
312        .rfind(|l| !l.is_empty())
313        .unwrap_or("failed")
314        .to_string()
315}
316
317/// Parse an ISO-8601 instant with a numeric offset (`2026-10-20T00:24:37+08:00`)
318/// into epoch ms plus the offset in seconds east of UTC.
319pub fn parse_instant(value: &str) -> Option<(i64, i64)> {
320    let (date, rest) = value.split_once('T')?;
321    let (time, offset) = match rest.find(['+', '-']) {
322        Some(at) => (&rest[..at], &rest[at..]),
323        None => (rest.trim_end_matches('Z'), "+00:00"),
324    };
325    let mut dp = date.split('-');
326    let y: i64 = dp.next()?.parse().ok()?;
327    let m: i64 = dp.next()?.parse().ok()?;
328    let d: i64 = dp.next()?.parse().ok()?;
329    let time = time.trim_end_matches('Z');
330    let mut tp = time.split(':');
331    let hh: i64 = tp.next().unwrap_or("0").parse().unwrap_or(0);
332    let mm: i64 = tp.next().unwrap_or("0").parse().unwrap_or(0);
333    let ss: i64 = tp
334        .next()
335        .unwrap_or("0")
336        .split('.')
337        .next()?
338        .parse()
339        .unwrap_or(0);
340    let (sign, off) = match offset.strip_prefix('-') {
341        Some(rest) => (-1, rest),
342        None => (1, offset.trim_start_matches('+')),
343    };
344    let mut op = off.split(':');
345    let oh: i64 = op.next().unwrap_or("0").parse().unwrap_or(0);
346    let om: i64 = op.next().unwrap_or("0").parse().unwrap_or(0);
347    let offset_secs = sign * (oh * 3600 + om * 60);
348    let local_secs = days_from_civil(y, m, d) * 86_400 + hh * 3600 + mm * 60 + ss;
349    Some(((local_secs - offset_secs) * 1000, offset_secs))
350}
351
352#[cfg(test)]
353mod tests {
354    use super::*;
355
356    #[test]
357    fn parses_the_servers_offset_instant() {
358        let (ms, off) = parse_instant("2026-10-20T00:24:37+08:00").unwrap();
359        assert_eq!(off, 8 * 3600);
360        assert_eq!(
361            ms,
362            (days_from_civil(2026, 10, 19) * 86_400 + 16 * 3600 + 24 * 60 + 37) * 1000
363        );
364        let (ms, off) = parse_instant("2026-10-20T00:24:37Z").unwrap();
365        assert_eq!(off, 0);
366        assert_eq!(
367            ms,
368            (days_from_civil(2026, 10, 20) * 86_400 + 24 * 60 + 37) * 1000
369        );
370    }
371
372    #[test]
373    fn reads_the_four_windows_and_the_tier() {
374        let raw = r#"{"items":[{"edition":"personal","tier":"medium","subscribed":true,
375            "periods":[
376              {"label":"5h","used":2975.67,"total":10000,"percent":29.75,"reset_at":"2026-10-20T00:24:37+08:00"},
377              {"label":"daily","used":100.0,"total":50000,"percent":0.2,"reset_at":"2026-10-20T00:00:00+08:00"},
378              {"label":"weekly","used":2975.67,"total":35000,"percent":8.5,"reset_at":"2026-10-26T00:00:00+08:00"},
379              {"label":"monthly","used":2975.67,"total":100000,"percent":2.98,"reset_at":"2026-11-19T23:59:59+08:00"}]}]}"#;
380        let usage = parse_usage(raw, "agent-plan_cn-beijing_medium").unwrap();
381        assert_eq!(usage.tier.as_deref(), Some("medium"));
382        assert_eq!(usage.five_hour.quota, 10_000.0);
383        assert!(usage.daily.is_some());
384        assert!(usage.five_hour.reset_ms().is_some());
385    }
386
387    #[test]
388    fn a_text_only_plan_has_no_daily_window() {
389        let raw = r#"{"items":[{"edition":"personal","tier":"small","subscribed":true,
390            "periods":[
391              {"label":"5h","used":1,"total":2000,"percent":0.05,"reset_at":"2026-10-20T00:24:37+08:00"},
392              {"label":"weekly","used":1,"total":7000,"percent":0.01,"reset_at":"2026-10-26T00:00:00+08:00"},
393              {"label":"monthly","used":1,"total":20000,"percent":0.005,"reset_at":"2026-11-19T23:59:59+08:00"}]}]}"#;
394        let usage = parse_usage(raw, "p").unwrap();
395        assert!(usage.daily.is_none());
396    }
397
398    #[test]
399    fn an_unsubscribed_item_is_an_error() {
400        let raw =
401            r#"{"items":[{"edition":"personal","tier":"medium","subscribed":false,"periods":[]}]}"#;
402        assert!(matches!(parse_usage(raw, "p"), Err(Unavailable::Failed(_))));
403    }
404}