brokk-mj-controller 2.29.0

//! Grok Build subscription usage querying through its ACP billing extension.
//!
//! Grok Build publishes no HTTP quota endpoint. Its own `/usage` view polls an
//! ACP ext request, `_x.ai/billing`, served by `grok agent stdio`. Keep the
//! JSON-RPC exchange isolated from the UI so the response parser stays
//! testable against captured payloads without spawning `grok`.

use std::collections::HashMap;
use std::fmt;
use std::path::{Path, PathBuf};
use std::process::Stdio;
use std::time::Duration;

use serde_json::{Value, json};
use tokio::io::{AsyncWriteExt, BufReader};
use tokio::process::{Child, ChildStdin, ChildStdout, Command};

/// A cold agent start plus one round trip to the billing backend. Generous
/// because the alternative — reporting a timeout on a slow machine — reads to
/// the user as a broken account.
const REQUEST_TIMEOUT: Duration = Duration::from_secs(30);
const MAX_RESPONSE_BYTES: usize = 256 * 1024;

/// Ext method Grok Build serves for billing. The `_` prefix is the ACP ext
/// framing, which is how the method appears on the wire.
const BILLING_METHOD: &str = "_x.ai/billing";

#[derive(Debug, Clone, PartialEq)]
pub struct GrokUsageReport {
    /// Window name for the quota row: `Week` or `Month`, matching the period
    /// the account is billed on.
    pub period_label: String,
    /// Included-credit usage as a percentage of the allowance, 0.0 to 100.0.
    pub used_percent: f64,
    /// Epoch seconds when the current period ends.
    pub resets_at: Option<i64>,
}

impl GrokUsageReport {
    /// Share of the allowance still available, rounded for the quota bar.
    pub fn remaining_percent(&self) -> u8 {
        let remaining = (100.0 - self.used_percent).round();
        remaining.clamp(0.0, 100.0) as u8
    }
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum GrokUsageError {
    NotInstalled,
    Launch(String),
    TimedOut,
    /// The agent exited before answering. Grok Build tears the process down
    /// when its stdin closes, so this also catches an early hangup.
    Closed,
    Protocol(String),
    /// The agent answered the billing request with a JSON-RPC error.
    Rejected(String),
    /// The agent answered, but with no billing configuration in it.
    NoData,
}

impl fmt::Display for GrokUsageError {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            Self::NotInstalled => write!(f, "Grok Build executable not found"),
            Self::Launch(error) => write!(f, "launch grok agent stdio: {error}"),
            Self::TimedOut => write!(f, "grok usage request timed out"),
            Self::Closed => write!(
                f,
                "grok agent stdio exited before answering the usage request"
            ),
            Self::Protocol(detail) => write!(f, "grok usage response was unreadable: {detail}"),
            Self::Rejected(detail) => write!(f, "grok refused the usage request: {detail}"),
            Self::NoData => write!(f, "grok reported no billing configuration"),
        }
    }
}

/// Ask a profile's Grok Build for its current billing period.
pub async fn query(
    home: PathBuf,
    cwd: PathBuf,
    env: HashMap<String, String>,
) -> Result<GrokUsageReport, GrokUsageError> {
    let mut session = GrokUsageSession::spawn(&home, cwd, env)?;
    // The exchange is bounded as a whole: a stalled agent must not hold the
    // quota refresh open, and its child must still be reaped.
    let result = tokio::time::timeout(REQUEST_TIMEOUT, session.exchange()).await;
    session.shutdown().await;
    result.map_err(|_| GrokUsageError::TimedOut)?
}

/// Executable candidates, in the order Hel tries them.
///
/// The installed CLI on `PATH` comes first, then the copy the official
/// installer leaves inside the profile home.
fn grok_programs(home: &Path) -> Vec<PathBuf> {
    let name = if cfg!(windows) { "grok.exe" } else { "grok" };
    vec![PathBuf::from(name), home.join("bin").join(name)]
}

struct GrokUsageSession {
    child: Child,
    /// Held for the whole exchange on purpose: Grok Build stops when its stdin
    /// closes, so dropping this early kills the agent before it answers.
    stdin: ChildStdin,
    stdout: BufReader<ChildStdout>,
}

impl GrokUsageSession {
    fn spawn(
        home: &Path,
        cwd: PathBuf,
        env: HashMap<String, String>,
    ) -> Result<Self, GrokUsageError> {
        let programs = grok_programs(home);
        let mut child = None;
        for (index, program) in programs.iter().enumerate() {
            let mut command = Command::new(program);
            command
                .args(["agent", "stdio"])
                .current_dir(&cwd)
                .envs(&env)
                .stdin(Stdio::piped())
                .stdout(Stdio::piped())
                .stderr(Stdio::null())
                .kill_on_drop(true);
            match command.spawn() {
                Ok(spawned) => {
                    child = Some(spawned);
                    break;
                }
                Err(error)
                    if error.kind() == std::io::ErrorKind::NotFound
                        && index + 1 < programs.len() => {}
                Err(error) if error.kind() == std::io::ErrorKind::NotFound => {
                    return Err(GrokUsageError::NotInstalled);
                }
                Err(error) => return Err(GrokUsageError::Launch(error.to_string())),
            }
        }
        let mut child = child.ok_or(GrokUsageError::NotInstalled)?;
        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| GrokUsageError::Protocol("agent stdin unavailable".into()))?;
        let stdout = child
            .stdout
            .take()
            .ok_or_else(|| GrokUsageError::Protocol("agent stdout unavailable".into()))?;
        Ok(Self {
            child,
            stdin,
            stdout: BufReader::new(stdout),
        })
    }

    async fn exchange(&mut self) -> Result<GrokUsageReport, GrokUsageError> {
        self.write_message(&json!({
            "jsonrpc": "2.0",
            "id": 1,
            "method": "initialize",
            "params": {"protocolVersion": 1, "clientCapabilities": {}},
        }))
        .await?;
        self.read_result(1).await?;

        self.write_message(&json!({
            "jsonrpc": "2.0",
            "id": 2,
            "method": BILLING_METHOD,
            "params": {},
        }))
        .await?;
        let billing = self.read_result(2).await?;
        parse(&billing)
    }

    async fn write_message(&mut self, message: &Value) -> Result<(), GrokUsageError> {
        let mut encoded = serde_json::to_vec(message)
            .map_err(|error| GrokUsageError::Protocol(error.to_string()))?;
        encoded.push(b'\n');
        self.stdin
            .write_all(&encoded)
            .await
            .map_err(|_| GrokUsageError::Closed)?;
        self.stdin.flush().await.map_err(|_| GrokUsageError::Closed)
    }

    /// Read frames until the answer to `expected_id` arrives. Everything else
    /// on the stream is the agent's own startup notifications.
    async fn read_result(&mut self, expected_id: u64) -> Result<Value, GrokUsageError> {
        loop {
            let Some(frame) = read_bounded_frame(&mut self.stdout).await? else {
                return Err(GrokUsageError::Closed);
            };
            let Ok(message) = serde_json::from_slice::<Value>(&frame) else {
                // A login-shell banner or other non-JSON noise is not fatal on
                // its own; the answer may still be coming.
                continue;
            };
            if message.get("id").and_then(Value::as_u64) != Some(expected_id) {
                continue;
            }
            if let Some(error) = message.get("error") {
                let detail = error
                    .get("message")
                    .and_then(Value::as_str)
                    .unwrap_or("no message")
                    .to_owned();
                return Err(GrokUsageError::Rejected(detail));
            }
            return message
                .get("result")
                .cloned()
                .ok_or_else(|| GrokUsageError::Protocol("response carries no result".into()));
        }
    }

    async fn shutdown(mut self) {
        drop(self.stdin);
        if let Err(error) = self.child.start_kill()
            && error.kind() != std::io::ErrorKind::NotFound
        {
            tracing::warn!(%error, "could not stop the Grok usage process");
        }
        if let Err(error) = self.child.wait().await {
            tracing::warn!(%error, "could not reap the Grok usage process");
        }
    }
}

async fn read_bounded_frame<R>(reader: &mut R) -> Result<Option<Vec<u8>>, GrokUsageError>
where
    R: tokio::io::AsyncBufRead + Unpin,
{
    use mj_core::bounded_frame::{BoundedFrame, BoundedFrameError};
    match mj_core::bounded_frame::read_bounded_frame(reader, MAX_RESPONSE_BYTES).await {
        Ok(BoundedFrame::Line(frame) | BoundedFrame::Truncated(frame)) => Ok(Some(frame)),
        Ok(BoundedFrame::End) => Ok(None),
        Err(BoundedFrameError::TooLarge) => Err(GrokUsageError::Protocol(
            "response frame is too large".into(),
        )),
        Err(BoundedFrameError::Io(error)) => Err(GrokUsageError::Protocol(error.to_string())),
    }
}

/// Project Grok Build's billing response into one quota window.
///
/// The top level is snake_case, but `config` is camelCase, matching the two
/// serde shapes in grok-build. proto3 JSON omits zero scalars, so an absent
/// `creditUsagePercent` means zero usage rather than missing data. The
/// deprecated `monthlyLimit`/`used`/`billingPeriod*` fields are read only when
/// their replacements are absent, exactly as grok's own struct documents.
pub fn parse(result: &Value) -> Result<GrokUsageReport, GrokUsageError> {
    let config = result.get("config").ok_or(GrokUsageError::NoData)?;
    if config.is_null() {
        return Err(GrokUsageError::NoData);
    }
    let period_type = config
        .pointer("/currentPeriod/type")
        .and_then(Value::as_str)
        .unwrap_or_default();
    let deprecated_only = config.get("currentPeriod").is_none();
    let period_label = if period_type.ends_with("WEEKLY") {
        "Week"
    } else if period_type.ends_with("MONTHLY") || deprecated_only {
        // The legacy shape only ever described a monthly credit budget.
        "Month"
    } else {
        "Period"
    };
    let used_percent = config
        .get("creditUsagePercent")
        .and_then(Value::as_f64)
        .or_else(|| {
            // Deprecated fallback: derive the share from the cent amounts.
            let limit = cents(config.get("monthlyLimit"))?;
            (limit > 0)
                .then(|| cents(config.get("used")).unwrap_or(0) as f64 * 100.0 / limit as f64)
        })
        .unwrap_or(0.0)
        .clamp(0.0, 100.0);
    let resets_at = config
        .pointer("/currentPeriod/end")
        .or_else(|| config.get("billingPeriodEnd"))
        .and_then(Value::as_str)
        .and_then(epoch_seconds);
    Ok(GrokUsageReport {
        period_label: period_label.to_owned(),
        used_percent,
        resets_at,
    })
}

/// A proto3 `Cent` message. `{}` is a valid zero, and so is an absent field.
fn cents(value: Option<&Value>) -> Option<i64> {
    let value = value?;
    Some(value.get("val").and_then(Value::as_i64).unwrap_or(0))
}

fn epoch_seconds(timestamp: &str) -> Option<i64> {
    chrono::DateTime::parse_from_rfc3339(timestamp)
        .ok()
        .map(|moment| moment.timestamp())
}

#[cfg(test)]
mod tests {
    use super::*;

    fn report(result: serde_json::Value) -> GrokUsageReport {
        parse(&result).unwrap()
    }

    /// proto3 JSON omits zero scalars, which is exactly what an account at
    /// zero usage looks like on the wire. Captured from a real account.
    #[test]
    fn omitted_zero_fields_read_as_zero_rather_than_missing_data() {
        let usage = report(json!({
            "config": {
                "currentPeriod": {
                    "type": "USAGE_PERIOD_TYPE_WEEKLY",
                    "start": "2026-08-11T05:22:07.661951+00:00",
                    "end": "2026-08-18T05:22:07.661951+00:00",
                },
                "onDemandCap": {},
                "onDemandUsed": {"val": 0},
                "prepaidBalance": {},
                "isUnifiedBillingUser": true,
                "billingPeriodStart": "2026-08-11T05:22:07.661951+00:00",
                "billingPeriodEnd": "2026-08-18T05:22:07.661951+00:00",
            },
            "subscription_tier": "X Premium+",
        }));

        assert_eq!(usage.used_percent, 0.0);
        assert_eq!(usage.remaining_percent(), 100);
        assert_eq!(usage.period_label, "Week");
        assert_eq!(cents(Some(&json!({}))), Some(0));
    }

    #[test]
    fn a_response_without_a_configuration_is_not_a_zero_reading() {
        assert_eq!(parse(&json!({})), Err(GrokUsageError::NoData));
        assert_eq!(parse(&json!({"config": null})), Err(GrokUsageError::NoData));
    }

    #[test]
    fn a_zero_allowance_does_not_divide_by_zero() {
        let usage = report(json!({
            "config": {"monthlyLimit": {"val": 0}, "used": {"val": 100}},
        }));

        assert_eq!(usage.used_percent, 0.0);
    }

    #[cfg(unix)]
    fn fake_grok(directory: &Path, script: &str) -> PathBuf {
        mj_core::test_hooks::install_fake_command(directory, "grok", script);
        directory.join("grok")
    }

    #[cfg(unix)]
    #[tokio::test]
    async fn an_agent_that_exits_without_answering_is_reported_not_awaited() {
        let directory = tempfile::tempdir().unwrap();
        let _executable = fake_grok(directory.path(), "#!/bin/sh\nexit 0\n");

        let error = query(
            directory.path().to_path_buf(),
            directory.path().to_path_buf(),
            HashMap::from([(
                "PATH".to_owned(),
                directory.path().to_string_lossy().into_owned(),
            )]),
        )
        .await
        .unwrap_err();

        assert_eq!(error, GrokUsageError::Closed);
        assert!(
            format!("{error}").contains("grok agent stdio exited"),
            "{error}"
        );
    }
}