From d3794bd77880081c26f9f2d3892a52aa7523f29f Mon Sep 17 00:00:00 2001 From: Factory Agent Date: Wed, 21 Oct 2026 12:25:36 +0700 Subject: [PATCH 1/3] Add Grok CLI ACP billing fallback When the bearer billing path fails for the ambient Grok CLI login, ask the installed grok agent stdio over ACP for _x.ai/billing. Isolated token accounts, saved account rows and OAuth principals never reach it. --- rust/src/providers/grok/acp.rs | 361 ++++++++++++ rust/src/providers/grok/acp_tests.rs | 522 ++++++++++++++++++ rust/src/providers/grok/cli_fallback_tests.rs | 223 ++++++++ rust/src/providers/grok/credits_proxy.rs | 4 +- rust/src/providers/grok/mod.rs | 93 +++- 5 files changed, 1199 insertions(+), 4 deletions(-) create mode 100644 rust/src/providers/grok/acp.rs create mode 100644 rust/src/providers/grok/acp_tests.rs create mode 100644 rust/src/providers/grok/cli_fallback_tests.rs diff --git a/rust/src/providers/grok/acp.rs b/rust/src/providers/grok/acp.rs new file mode 100644 index 00000000000..3fb3095d0d1 --- /dev/null +++ b/rust/src/providers/grok/acp.rs @@ -0,0 +1,361 @@ +//! Grok CLI billing over the ACP extension method `_x.ai/billing` (upstream #4595). +//! +//! `grok agent stdio` speaks newline-delimited JSON-RPC 2.0 on stdin and stdout. +//! CodexBar sends `initialize`, then `_x.ai/billing` with empty params. The CLI +//! answers from its own login, so no token from CodexBar is sent. +//! +//! This is a fallback only. `mod.rs` calls it after the bearer billing path has +//! failed, and only for the ambient Grok CLI login. Isolated token accounts and +//! saved account rows never reach it. + +use std::io; +use std::process::Stdio; +use std::time::Duration; + +use chrono::{DateTime, Utc}; +use serde::Deserialize; +use serde_json::{Value, json}; +use tokio::io::{AsyncBufRead, AsyncBufReadExt, AsyncWrite, AsyncWriteExt, BufReader}; +use tokio::process::{Child, Command}; + +use crate::core::ProviderError; + +use super::billing::GrokBillingSnapshot; +use super::credits_proxy::{ + BearerBilling, parse_credits_response, parse_timestamp, window_minutes, +}; +use super::grok_plan_display_name; + +const BILLING_METHOD: &str = "_x.ai/billing"; +const INITIALIZE_ID: u64 = 1; +const BILLING_ID: u64 = 2; +const METHOD_NOT_FOUND: i64 = -32601; +/// Upper bound on one stdout line. Billing answers are small; anything larger +/// is not a reply CodexBar should buffer. +const MAX_LINE_BYTES: usize = 256 * 1024; +/// How long a killed child gets to be reaped before the caller moves on. +const EXIT_GRACE: Duration = Duration::from_secs(1); +const USAGE_UNAVAILABLE: &str = "Grok CLI billing returned no usage, reset or prepaid balance"; + +#[cfg(windows)] +const CREATE_NO_WINDOW: u32 = 0x0800_0000; + +/// Per-request budgets, matching upstream `GrokRPCClient` (4s for initialize, +/// 3s for billing). +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(super) struct AcpTimeouts { + pub(super) initialize: Duration, + pub(super) billing: Duration, +} + +impl Default for AcpTimeouts { + fn default() -> Self { + Self { + initialize: Duration::from_secs(4), + billing: Duration::from_secs(3), + } + } +} + +/// Runs `program args...`, asks it for billing over ACP, and returns the answer. +/// +/// The child has no console window, its stderr is discarded, and it is killed +/// before this returns on every path: success, error or timeout. +pub(super) async fn fetch_billing( + program: &str, + args: &[&str], + timeouts: AcpTimeouts, +) -> Result { + let mut child = spawn_child(program, args)?; + let outcome = run_child(&mut child, timeouts, Utc::now()).await; + stop(&mut child).await; + outcome +} + +fn spawn_child(program: &str, args: &[&str]) -> Result { + let mut command = Command::new(program); + command + .args(args) + .stdin(Stdio::piped()) + .stdout(Stdio::piped()) + .stderr(Stdio::null()) + .kill_on_drop(true); + #[cfg(windows)] + command.creation_flags(CREATE_NO_WINDOW); + + command.spawn().map_err(|error| { + if error.kind() == io::ErrorKind::NotFound { + ProviderError::NotInstalled( + "Grok CLI not found. Install it, then run `grok login`.".to_string(), + ) + } else { + ProviderError::Other("Grok CLI failed to start".to_string()) + } + }) +} + +async fn run_child( + child: &mut Child, + timeouts: AcpTimeouts, + now: DateTime, +) -> Result { + let (Some(mut stdin), Some(stdout)) = (child.stdin.take(), child.stdout.take()) else { + return Err(ProviderError::Other( + "Grok CLI did not open its ACP pipes".to_string(), + )); + }; + let mut stdout = BufReader::new(stdout); + billing_over_pipes(&mut stdout, &mut stdin, timeouts, now).await +} + +/// `grok agent stdio` is a long-running server, so every use stops it. A timeout +/// takes the same path: the child is killed, never left running. +async fn stop(child: &mut Child) { + child.start_kill().ok(); + tokio::time::timeout(EXIT_GRACE, child.wait()).await.ok(); +} + +/// The handshake and billing request over an already-open pipe pair. +pub(super) async fn billing_over_pipes( + reader: &mut R, + writer: &mut W, + timeouts: AcpTimeouts, + now: DateTime, +) -> Result +where + R: AsyncBufRead + Unpin, + W: AsyncWrite + Unpin, +{ + request( + reader, + writer, + INITIALIZE_ID, + "initialize", + initialize_params(), + timeouts.initialize, + ) + .await?; + let result = request( + reader, + writer, + BILLING_ID, + BILLING_METHOD, + json!({}), + timeouts.billing, + ) + .await?; + parse_billing_result(&result, now) +} + +fn initialize_params() -> Value { + json!({ + "protocolVersion": "1", + "clientCapabilities": { + "fs": { "readTextFile": false, "writeTextFile": false }, + "terminal": false, + }, + }) +} + +async fn request( + reader: &mut R, + writer: &mut W, + id: u64, + method: &str, + params: Value, + budget: Duration, +) -> Result +where + R: AsyncBufRead + Unpin, + W: AsyncWrite + Unpin, +{ + let mut line = serde_json::to_vec(&json!({ + "jsonrpc": "2.0", + "id": id, + "method": method, + "params": params, + })) + .map_err(|_| ProviderError::Other("Could not encode the Grok CLI request".to_string()))?; + line.push(b'\n'); + match tokio::time::timeout(budget, exchange(reader, writer, &line, id)).await { + Ok(result) => result, + Err(_elapsed) => Err(ProviderError::Timeout), + } +} + +async fn exchange( + reader: &mut R, + writer: &mut W, + line: &[u8], + id: u64, +) -> Result +where + R: AsyncBufRead + Unpin, + W: AsyncWrite + Unpin, +{ + writer.write_all(line).await.map_err(pipe_error)?; + writer.flush().await.map_err(pipe_error)?; + loop { + let Some(bytes) = read_bounded_line(reader).await.map_err(pipe_error)? else { + return Err(ProviderError::Other( + "Grok CLI closed its ACP output before replying".to_string(), + )); + }; + // Lines that are not JSON, and notifications, are not replies. + let Ok(message) = serde_json::from_slice::(&bytes) else { + continue; + }; + if let Some(reply) = reply_for(&message, id) { + return reply; + } + } +} + +fn pipe_error(error: io::Error) -> ProviderError { + match error.kind() { + io::ErrorKind::InvalidData => { + ProviderError::Other("Grok CLI sent an oversized ACP message".to_string()) + } + kind => ProviderError::Other(format!("Grok CLI ACP pipe failed: {kind:?}")), + } +} + +/// Reads one `\n`-terminated line, without the size of a line ever exceeding +/// [`MAX_LINE_BYTES`]. `None` means the stream ended with nothing buffered. +async fn read_bounded_line(reader: &mut R) -> io::Result>> { + let mut line = Vec::new(); + loop { + let buffered = reader.fill_buf().await?; + if buffered.is_empty() { + return Ok((!line.is_empty()).then_some(line)); + } + let newline = buffered.iter().position(|byte| *byte == b'\n'); + let take = newline.map_or(buffered.len(), |index| index + 1); + if line.len() + take > MAX_LINE_BYTES { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "ACP line exceeds the size limit", + )); + } + line.extend_from_slice(&buffered[..take]); + reader.consume(take); + if newline.is_some() { + return Ok(Some(line)); + } + } +} + +/// The reply to request `id`, or `None` for notifications and replies to other +/// requests. +fn reply_for(message: &Value, id: u64) -> Option> { + let message_id = message.get("id")?.as_u64()?; + if message_id != id { + return None; + } + if let Some(error) = message.get("error").filter(|error| !error.is_null()) { + return Some(Err(rpc_error(error))); + } + match message.get("result") { + Some(result) => Some(Ok(result.clone())), + None => Some(Err(ProviderError::Parse( + "Grok CLI reply had no result".to_string(), + ))), + } +} + +fn rpc_error(error: &Value) -> ProviderError { + if error.get("code").and_then(Value::as_i64) == Some(METHOD_NOT_FOUND) { + return ProviderError::Other( + "This Grok CLI version does not provide billing over ACP".to_string(), + ); + } + let message = error + .get("message") + .and_then(Value::as_str) + .unwrap_or("unknown JSON-RPC error"); + let lowered = message.to_ascii_lowercase(); + if lowered.contains("authentication required") || lowered.contains("grok login") { + return ProviderError::AuthRequired; + } + ProviderError::Other(format!("Grok CLI billing request failed: {message}")) +} + +/// Decodes a `_x.ai/billing` result. The current shape is the nested `config` +/// credits payload, read by the shared credits parser. The older top-level +/// `monthlyLimit` / `usage.totalUsed` shape is still accepted. A top-level +/// `subscription_tier` names the plan in both cases. +pub(super) fn parse_billing_result( + result: &Value, + now: DateTime, +) -> Result { + let plan = grok_plan_display_name(result.get("subscription_tier").and_then(Value::as_str)); + if let Ok(body) = serde_json::to_vec(result) + && let Ok(mut credits) = parse_credits_response(&body, now) + && (credits.billing.used_percent.is_some() || credits.prepaid_balance_usd.is_some()) + { + credits.subscription_tier = plan.or(credits.subscription_tier); + return Ok(credits); + } + if let Some(billing) = legacy_billing(result, now) { + return Ok(BearerBilling { + billing, + subscription_tier: plan, + prepaid_balance_usd: None, + }); + } + Err(ProviderError::Parse(USAGE_UNAVAILABLE.to_string())) +} + +/// Legacy top-level billing shape. Monetary values are `{ "val": }`. +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct LegacyBilling { + billing_cycle: Option, + monthly_limit: Option, + usage: Option, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct LegacyCycle { + billing_period_start: Option, + billing_period_end: Option, +} + +#[derive(Deserialize)] +#[serde(rename_all = "camelCase")] +struct LegacyUsage { + total_used: Option, +} + +#[derive(Deserialize)] +struct LegacyCent { + val: Option, +} + +/// Usage needs a positive limit and a used amount. Without both there is no +/// percent, so the payload is not accepted as legacy billing. +fn legacy_billing(result: &Value, now: DateTime) -> Option { + let legacy = LegacyBilling::deserialize(result).ok()?; + let cycle = legacy.billing_cycle.as_ref(); + let resets_at = parse_timestamp(cycle.and_then(|cycle| cycle.billing_period_end.as_deref())); + let minutes = window_minutes( + cycle.and_then(|cycle| cycle.billing_period_start.as_deref()), + resets_at, + now, + ); + let limit = legacy.monthly_limit?.val.filter(|limit| *limit > 0.0)?; + let used = legacy.usage?.total_used?.val?; + Some(GrokBillingSnapshot { + used_percent: Some((used / limit * 100.0).clamp(0.0, 100.0)), + used_percent_is_wire_published: true, + used_percent_is_implicit_zero: false, + resets_at, + window_minutes: minutes, + product_usage: Vec::new(), + }) +} + +#[cfg(test)] +#[path = "acp_tests.rs"] +mod tests; diff --git a/rust/src/providers/grok/acp_tests.rs b/rust/src/providers/grok/acp_tests.rs new file mode 100644 index 00000000000..9846e28b585 --- /dev/null +++ b/rust/src/providers/grok/acp_tests.rs @@ -0,0 +1,522 @@ +use super::*; +use chrono::TimeZone; +use std::time::Instant; +use tokio::io::{AsyncBufReadExt, duplex}; + +fn now() -> DateTime { + Utc.with_ymd_and_hms(2026, 1, 15, 0, 0, 0).unwrap() +} + +fn short_timeouts() -> AcpTimeouts { + AcpTimeouts { + initialize: Duration::from_millis(200), + billing: Duration::from_millis(200), + } +} + +/// Nested `config` payload as the current `_x.ai/billing` result carries it. +fn credits_result(percent: f64) -> Value { + json!({ + "config": { + "creditUsagePercent": percent, + "currentPeriod": { + "start": "2026-01-01T00:00:00Z", + "end": "2026-02-01T00:00:00Z", + }, + }, + }) +} + +#[test] +fn current_result_maps_usage_reset_window_and_plan() { + let mut result = credits_result(37.5); + result["subscription_tier"] = json!("heavy"); + + let bearer = parse_billing_result(&result, now()).unwrap(); + + assert_eq!(bearer.billing.used_percent, Some(37.5)); + assert!(bearer.billing.used_percent_is_wire_published); + assert_eq!( + bearer.billing.resets_at, + Utc.with_ymd_and_hms(2026, 2, 1, 0, 0, 0).single() + ); + // 2026-01-01 to 2026-02-01 is 31 days. + assert_eq!(bearer.billing.window_minutes, Some(31 * 24 * 60)); + assert_eq!(bearer.subscription_tier.as_deref(), Some("SuperGrok Heavy")); + assert_eq!(bearer.prepaid_balance_usd, None); +} + +#[test] +fn top_level_tier_overrides_the_config_tier_and_config_tier_is_the_fallback() { + let config_only = json!({ + "config": { "creditUsagePercent": 5.0, "subscriptionTier": "SuperGrok" }, + }); + assert_eq!( + parse_billing_result(&config_only, now()) + .unwrap() + .subscription_tier + .as_deref(), + Some("SuperGrok") + ); + + let mut both = config_only.clone(); + both["subscription_tier"] = json!("heavy"); + assert_eq!( + parse_billing_result(&both, now()) + .unwrap() + .subscription_tier + .as_deref(), + Some("SuperGrok Heavy") + ); +} + +#[test] +fn prepaid_balance_alone_is_a_valid_answer_without_a_quota_percent() { + let result = json!({ + "config": { "prepaidBalance": { "val": 1250 } }, + }); + + let bearer = parse_billing_result(&result, now()).unwrap(); + + assert_eq!(bearer.billing.used_percent, None); + assert_eq!(bearer.billing.resets_at, None); + assert_eq!(bearer.prepaid_balance_usd, Some(12.5)); +} + +#[test] +fn prepaid_balance_reads_string_cents_and_an_empty_wallet_as_zero() { + let string_cents = json!({ "config": { "prepaidBalance": { "val": "250" } } }); + assert_eq!( + parse_billing_result(&string_cents, now()) + .unwrap() + .prepaid_balance_usd, + Some(2.5) + ); + + let empty_wallet = json!({ "config": { "prepaidBalance": {} } }); + assert_eq!( + parse_billing_result(&empty_wallet, now()) + .unwrap() + .prepaid_balance_usd, + Some(0.0) + ); +} + +#[test] +fn quota_and_prepaid_balance_arrive_together() { + let mut result = credits_result(10.0); + result["config"]["prepaidBalance"] = json!({ "val": 500 }); + + let bearer = parse_billing_result(&result, now()).unwrap(); + + assert_eq!(bearer.billing.used_percent, Some(10.0)); + assert_eq!(bearer.prepaid_balance_usd, Some(5.0)); +} + +#[test] +fn reset_time_alone_is_not_usage() { + let result = json!({ + "config": { + "currentPeriod": { "start": "2026-01-01T00:00:00Z", "end": "2026-02-01T00:00:00Z" }, + }, + }); + + assert!(matches!( + parse_billing_result(&result, now()), + Err(ProviderError::Parse(_)) + )); +} + +#[test] +fn legacy_top_level_shape_still_maps_usage_and_reset() { + let result = json!({ + "billingCycle": { + "billingPeriodStart": "2026-01-01T00:00:00Z", + "billingPeriodEnd": "2026-02-01T00:00:00Z", + }, + "monthlyLimit": { "val": 100 }, + "usage": { "totalUsed": { "val": 25 } }, + "subscription_tier": "SuperGrok", + }); + + let bearer = parse_billing_result(&result, now()).unwrap(); + + assert_eq!(bearer.billing.used_percent, Some(25.0)); + assert!(bearer.billing.used_percent_is_wire_published); + assert_eq!( + bearer.billing.resets_at, + Utc.with_ymd_and_hms(2026, 2, 1, 0, 0, 0).single() + ); + assert_eq!(bearer.billing.window_minutes, Some(31 * 24 * 60)); + assert_eq!(bearer.subscription_tier.as_deref(), Some("SuperGrok")); + assert_eq!(bearer.prepaid_balance_usd, None); +} + +#[test] +fn legacy_usage_is_clamped_to_the_limit() { + let result = json!({ + "monthlyLimit": { "val": 100 }, + "usage": { "totalUsed": { "val": 150 } }, + }); + + assert_eq!( + parse_billing_result(&result, now()) + .unwrap() + .billing + .used_percent, + Some(100.0) + ); +} + +#[test] +fn legacy_shape_without_a_positive_limit_is_unavailable() { + for result in [ + json!({ "monthlyLimit": { "val": 0 }, "usage": { "totalUsed": { "val": 5 } } }), + json!({ "usage": { "totalUsed": { "val": 5 } } }), + json!({ "monthlyLimit": { "val": 100 } }), + ] { + assert!( + matches!( + parse_billing_result(&result, now()), + Err(ProviderError::Parse(_)) + ), + "{result}" + ); + } +} + +#[test] +fn missing_or_null_fields_are_parse_errors_not_empty_snapshots() { + for result in [ + json!({}), + json!({ "config": {} }), + json!({ "config": null }), + json!({ "config": { "creditUsagePercent": "lots" } }), + Value::Null, + ] { + assert!( + matches!( + parse_billing_result(&result, now()), + Err(ProviderError::Parse(_)) + ), + "{result}" + ); + } +} + +#[test] +fn a_percent_without_a_reset_is_still_usage() { + let result = json!({ "config": { "creditUsagePercent": 10.0 } }); + + let bearer = parse_billing_result(&result, now()).unwrap(); + + assert_eq!(bearer.billing.used_percent, Some(10.0)); + assert_eq!(bearer.billing.resets_at, None); + assert_eq!(bearer.billing.window_minutes, None); +} + +#[test] +fn reply_for_matches_only_the_requested_id() { + // Notifications carry no id; replies to other requests carry another id. + assert!(reply_for(&json!({ "jsonrpc": "2.0", "method": "session/update" }), 2).is_none()); + assert!(reply_for(&json!({ "jsonrpc": "2.0", "id": 1, "result": {} }), 2).is_none()); + + let reply = reply_for( + &json!({ "jsonrpc": "2.0", "id": 2, "result": { "a": 1 } }), + 2, + ) + .expect("matching reply") + .expect("result is a success"); + assert_eq!(reply, json!({ "a": 1 })); +} + +#[test] +fn reply_for_treats_a_null_error_as_no_error_and_a_bare_id_as_malformed() { + let reply = reply_for( + &json!({ "jsonrpc": "2.0", "id": 2, "error": null, "result": {} }), + 2, + ) + .expect("matching reply") + .expect("null error is not an error"); + assert_eq!(reply, json!({})); + + assert!(matches!( + reply_for(&json!({ "jsonrpc": "2.0", "id": 2 }), 2), + Some(Err(ProviderError::Parse(_))) + )); +} + +#[test] +fn rpc_error_objects_map_to_auth_unsupported_and_other() { + assert!(matches!( + rpc_error(&json!({ + "code": -32000, + "message": "Authentication required to fetch billing data. Run 'grok login' to authenticate.", + })), + ProviderError::AuthRequired + )); + assert!(matches!( + rpc_error(&json!({ "code": -32601, "message": "Method not found" })), + ProviderError::Other(message) if message.contains("does not provide billing") + )); + assert!(matches!( + rpc_error(&json!({ "code": -32000, "message": "quota service down" })), + ProviderError::Other(message) if message.contains("quota service down") + )); + assert!(matches!( + rpc_error(&json!({ "code": -32000 })), + ProviderError::Other(message) if message.contains("unknown JSON-RPC error") + )); +} + +/// Two in-memory pipes: the client writes to the server and reads from it. +fn pipes() -> ( + tokio::io::DuplexStream, + BufReader, + tokio::io::DuplexStream, + tokio::io::DuplexStream, +) { + let (client_writer, server_reader) = duplex(64 * 1024); + let (server_writer, client_reader) = duplex(64 * 1024); + ( + client_writer, + BufReader::new(client_reader), + server_reader, + server_writer, + ) +} + +#[tokio::test] +async fn handshake_then_billing_skips_noise_and_decodes_the_answer() { + let (mut client_writer, mut client_reader, server_reader, mut server_writer) = pipes(); + let server = tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + let init: Value = serde_json::from_str(&lines.next_line().await.unwrap().unwrap()).unwrap(); + assert_eq!(init["method"], "initialize"); + assert_eq!(init["id"], 1); + assert_eq!(init["params"]["protocolVersion"], "1"); + assert_eq!(init["params"]["clientCapabilities"]["terminal"], false); + server_writer.write_all(b"not json at all\n").await.unwrap(); + server_writer + .write_all(b"{\"jsonrpc\":\"2.0\",\"method\":\"session/update\",\"params\":{}}\n") + .await + .unwrap(); + server_writer + .write_all(b"{\"jsonrpc\":\"2.0\",\"id\":1,\"result\":{}}\n") + .await + .unwrap(); + + let billing_line = lines.next_line().await.unwrap().unwrap(); + // The extension method must reach the wire literally, without a + // `\/` escape: grok does not unescape method names. + assert!(billing_line.contains("\"_x.ai/billing\""), "{billing_line}"); + let billing: Value = serde_json::from_str(&billing_line).unwrap(); + assert_eq!(billing["id"], 2); + assert_eq!(billing["method"], "_x.ai/billing"); + assert_eq!(billing["params"], json!({})); + let answer = json!({ "jsonrpc": "2.0", "id": 2, "result": credits_result(42.0) }); + server_writer + .write_all(format!("{answer}\n").as_bytes()) + .await + .unwrap(); + }); + + let bearer = billing_over_pipes( + &mut client_reader, + &mut client_writer, + short_timeouts(), + now(), + ) + .await + .unwrap(); + + assert_eq!(bearer.billing.used_percent, Some(42.0)); + server.await.unwrap(); +} + +#[tokio::test] +async fn silent_child_times_out_on_initialize() { + let (mut client_writer, mut client_reader, server_reader, server_writer) = pipes(); + // Keep the server ends open so the silence is a timeout, not end of stream. + let _held = (server_reader, server_writer); + + let started = Instant::now(); + let error = billing_over_pipes( + &mut client_reader, + &mut client_writer, + short_timeouts(), + now(), + ) + .await + .unwrap_err(); + + assert!(matches!(error, ProviderError::Timeout), "{error}"); + assert!(started.elapsed() < Duration::from_secs(5)); +} + +#[tokio::test] +async fn silent_child_times_out_on_billing_after_a_good_handshake() { + let (mut client_writer, mut client_reader, server_reader, mut server_writer) = pipes(); + let server = tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + let init = lines.next_line().await.unwrap().unwrap(); + let id = serde_json::from_str::(&init).unwrap()["id"].clone(); + let reply = json!({ "jsonrpc": "2.0", "id": id, "result": {} }); + server_writer + .write_all(format!("{reply}\n").as_bytes()) + .await + .unwrap(); + // Read the billing request, then never answer it. + let _billing = lines.next_line().await.unwrap(); + tokio::time::sleep(Duration::from_secs(30)).await; + }); + + let error = billing_over_pipes( + &mut client_reader, + &mut client_writer, + short_timeouts(), + now(), + ) + .await + .unwrap_err(); + + assert!(matches!(error, ProviderError::Timeout), "{error}"); + server.abort(); +} + +#[tokio::test] +async fn closed_output_before_a_reply_is_an_error() { + let (mut client_writer, mut client_reader, server_reader, server_writer) = pipes(); + // The request can be written, but the reply side closes. + let _held_reader = server_reader; + drop(server_writer); + + let error = billing_over_pipes( + &mut client_reader, + &mut client_writer, + short_timeouts(), + now(), + ) + .await + .unwrap_err(); + + assert!( + matches!(&error, ProviderError::Other(message) if message.contains("closed")), + "{error}" + ); +} + +#[tokio::test] +async fn an_auth_error_during_the_handshake_is_surfaced_as_auth_required() { + let (mut client_writer, mut client_reader, server_reader, mut server_writer) = pipes(); + let server = tokio::spawn(async move { + let mut lines = BufReader::new(server_reader).lines(); + let _init = lines.next_line().await.unwrap(); + let reply = json!({ + "jsonrpc": "2.0", + "id": 1, + "error": { "code": -32000, "message": "Authentication required" }, + }); + server_writer + .write_all(format!("{reply}\n").as_bytes()) + .await + .unwrap(); + }); + + let error = billing_over_pipes( + &mut client_reader, + &mut client_writer, + short_timeouts(), + now(), + ) + .await + .unwrap_err(); + + assert!(matches!(error, ProviderError::AuthRequired), "{error}"); + server.await.unwrap(); +} + +#[tokio::test] +async fn bounded_line_reader_returns_each_line_then_none_at_end_of_stream() { + let (mut writer, reader) = duplex(64 * 1024); + let mut reader = BufReader::new(reader); + writer.write_all(b"{\"a\":1}\nlast").await.unwrap(); + drop(writer); + + assert_eq!( + read_bounded_line(&mut reader).await.unwrap(), + Some(b"{\"a\":1}\n".to_vec()) + ); + // A final line without a newline is still delivered once. + assert_eq!( + read_bounded_line(&mut reader).await.unwrap(), + Some(b"last".to_vec()) + ); + assert_eq!(read_bounded_line(&mut reader).await.unwrap(), None); +} + +#[tokio::test] +async fn bounded_line_reader_rejects_a_line_over_the_limit() { + let (mut writer, reader) = duplex(64 * 1024); + let mut reader = BufReader::new(reader); + let feed = tokio::spawn(async move { + let oversized = vec![b'a'; MAX_LINE_BYTES + 1]; + writer.write_all(&oversized).await.ok(); + writer.write_all(b"\n").await.ok(); + }); + + let error = read_bounded_line(&mut reader).await.unwrap_err(); + + assert_eq!(error.kind(), io::ErrorKind::InvalidData); + drop(reader); + feed.await.ok(); +} + +/// A child that stays alive and never speaks ACP. +#[cfg(windows)] +const SILENT_CHILD: (&str, &[&str]) = ( + "powershell", + &[ + "-NoProfile", + "-NonInteractive", + "-Command", + "Start-Sleep -Seconds 120", + ], +); +#[cfg(not(windows))] +const SILENT_CHILD: (&str, &[&str]) = ("sleep", &["120"]); + +#[tokio::test] +async fn timed_out_child_is_stopped_and_reaped() { + let (program, args) = SILENT_CHILD; + let mut child = spawn_child(program, args).expect("silent child starts"); + + let outcome = run_child(&mut child, short_timeouts(), now()).await; + // `fetch_billing` runs exactly this sequence. + stop(&mut child).await; + + assert!( + matches!(outcome, Err(ProviderError::Timeout)), + "{outcome:?}" + ); + assert!( + child + .try_wait() + .expect("child status is readable") + .is_some(), + "the child must be reaped, not left running" + ); +} + +#[tokio::test] +async fn a_missing_grok_binary_is_reported_as_not_installed() { + let error = fetch_billing( + "codexbar-test-no-such-grok-binary", + &["agent", "stdio"], + short_timeouts(), + ) + .await + .unwrap_err(); + + assert!(matches!(error, ProviderError::NotInstalled(_)), "{error}"); +} diff --git a/rust/src/providers/grok/cli_fallback_tests.rs b/rust/src/providers/grok/cli_fallback_tests.rs new file mode 100644 index 00000000000..56823a6306a --- /dev/null +++ b/rust/src/providers/grok/cli_fallback_tests.rs @@ -0,0 +1,223 @@ +use super::*; +use std::sync::Arc; +use std::sync::atomic::{AtomicUsize, Ordering}; + +const CREDITS_OK_BODY: &str = r#"{"config":{"creditUsagePercent":12.5,"currentPeriod":{"start":"2026-01-01T00:00:00Z","end":"2099-02-01T00:00:00Z"}},"subscriptionTier":"SuperGrok"}"#; + +fn cli_answer() -> BearerBilling { + BearerBilling { + billing: GrokBillingSnapshot { + used_percent: Some(64.0), + used_percent_is_wire_published: true, + used_percent_is_implicit_zero: false, + resets_at: None, + window_minutes: None, + product_usage: Vec::new(), + }, + subscription_tier: Some("SuperGrok Heavy".to_string()), + prepaid_balance_usd: None, + } +} + +/// A fake CLI runner that counts how often it is asked. +fn counting_cli( + answer: fn() -> Result, +) -> ( + Arc, + impl Fn() -> Result + Send + Sync + 'static, +) { + let calls = Arc::new(AtomicUsize::new(0)); + let counter = Arc::clone(&calls); + (calls, move || { + counter.fetch_add(1, Ordering::SeqCst); + answer() + }) +} + +async fn provider_with_bearer_status( + server: &mut mockito::ServerGuard, + status: usize, +) -> GrokProvider { + server + .mock("GET", "/credits") + .with_status(status) + .create_async() + .await; + server + .mock("POST", "/billing") + .with_status(status) + .create_async() + .await; + GrokProvider::new() + .with_billing_endpoint_for_tests(format!("{}/billing", server.url())) + .with_credits_proxy_endpoint_for_tests(format!("{}/credits", server.url())) +} + +#[tokio::test] +async fn successful_bearer_billing_never_asks_the_cli() { + let mut server = mockito::Server::new_async().await; + let credits = server + .mock("GET", "/credits") + .with_status(200) + .with_body(CREDITS_OK_BODY) + .create_async() + .await; + let grpc = server + .mock("POST", "/billing") + .expect(0) + .create_async() + .await; + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = GrokProvider::new() + .with_billing_endpoint_for_tests(format!("{}/billing", server.url())) + .with_credits_proxy_endpoint_for_tests(format!("{}/credits", server.url())) + .with_cli_billing_for_tests(run); + + let bearer = provider + .fetch_bearer_or_cli_billing( + &GrokCredentials::from_bearer("token"), + GrokAuthKind::Cli, + &FetchContext::default(), + ) + .await + .unwrap(); + + credits.assert_async().await; + grpc.assert_async().await; + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert_eq!(bearer.billing.used_percent, Some(12.5)); +} + +#[tokio::test] +async fn failed_bearer_billing_falls_back_to_the_cli_once() { + let mut server = mockito::Server::new_async().await; + let provider = provider_with_bearer_status(&mut server, 500).await; + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = provider.with_cli_billing_for_tests(run); + + let bearer = provider + .fetch_bearer_or_cli_billing( + &GrokCredentials::from_bearer("token"), + GrokAuthKind::Cli, + &FetchContext::default(), + ) + .await + .unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(bearer.billing.used_percent, Some(64.0)); + assert_eq!(bearer.subscription_tier.as_deref(), Some("SuperGrok Heavy")); +} + +#[tokio::test] +async fn auth_required_on_the_bearer_path_also_falls_back_to_the_cli() { + let mut server = mockito::Server::new_async().await; + let provider = provider_with_bearer_status(&mut server, 401).await; + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = provider.with_cli_billing_for_tests(run); + + let bearer = provider + .fetch_bearer_or_cli_billing( + &GrokCredentials::from_bearer("token"), + GrokAuthKind::Cli, + &FetchContext::default(), + ) + .await + .unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(bearer.billing.used_percent, Some(64.0)); +} + +#[tokio::test] +async fn failing_cli_fallback_keeps_the_bearer_error_for_auto_mode() { + let mut server = mockito::Server::new_async().await; + let provider = provider_with_bearer_status(&mut server, 500).await; + // The CLI's own auth error must not replace the bearer error. + let (calls, run) = counting_cli(|| Err(ProviderError::AuthRequired)); + let provider = provider.with_cli_billing_for_tests(run); + + let error = provider + .fetch_bearer_or_cli_billing( + &GrokCredentials::from_bearer("token"), + GrokAuthKind::Cli, + &FetchContext::default(), + ) + .await + .unwrap_err(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!( + matches!(&error, ProviderError::Other(message) if message.contains("status 500")), + "{error}" + ); +} + +#[tokio::test] +async fn isolated_token_account_never_reaches_the_cli() { + let mut server = mockito::Server::new_async().await; + let provider = provider_with_bearer_status(&mut server, 500).await; + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = provider.with_cli_billing_for_tests(run); + let ctx = FetchContext { + token_account_isolated: true, + ..FetchContext::default() + }; + + let error = provider + .fetch_bearer_or_cli_billing( + &GrokCredentials::from_bearer("token"), + GrokAuthKind::Cli, + &ctx, + ) + .await + .unwrap_err(); + + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert!(matches!(error, ProviderError::Other(_)), "{error}"); +} + +#[tokio::test] +async fn oauth_principal_never_reaches_the_cli() { + let mut server = mockito::Server::new_async().await; + let provider = provider_with_bearer_status(&mut server, 401).await; + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = provider.with_cli_billing_for_tests(run); + + let error = provider + .fetch_bearer_or_cli_billing( + &GrokCredentials::from_bearer("token"), + GrokAuthKind::OAuth, + &FetchContext::default(), + ) + .await + .unwrap_err(); + + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert!(matches!(error, ProviderError::AuthRequired), "{error}"); +} + +#[test] +fn saved_account_rows_are_isolated_from_the_ambient_cli_login() { + assert!(saved_account_usage_context().token_account_isolated); +} + +#[tokio::test] +async fn saved_account_row_failure_never_reaches_the_ambient_cli() { + // A saved row's auth text may name another account than the one switched + // into ~/.grok, so a failed row must not borrow the ambient login over ACP. + let mut server = mockito::Server::new_async().await; + let provider = provider_with_bearer_status(&mut server, 500).await; + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = provider.with_cli_billing_for_tests(run); + let saved_row = + r#"{"https://accounts.x.ai/sign-in": {"key": "saved-cli-token", "auth_mode": "session"}}"#; + + let error = provider + .fetch_usage_from_auth_json(saved_row) + .await + .unwrap_err(); + + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert!(matches!(error, ProviderError::Other(_)), "{error}"); +} diff --git a/rust/src/providers/grok/credits_proxy.rs b/rust/src/providers/grok/credits_proxy.rs index 18b7a2b2fd8..e91c085c739 100644 --- a/rust/src/providers/grok/credits_proxy.rs +++ b/rust/src/providers/grok/credits_proxy.rs @@ -113,7 +113,7 @@ fn usd_from_exact_cents(cents: i64) -> Option { .then(|| value / 100.0) } -fn parse_timestamp(raw: Option<&str>) -> Option> { +pub(super) fn parse_timestamp(raw: Option<&str>) -> Option> { DateTime::parse_from_rfc3339(raw?) .ok() .map(|parsed| parsed.with_timezone(&Utc)) @@ -123,7 +123,7 @@ fn parse_timestamp(raw: Option<&str>) -> Option> { /// is not in the future, so a reset timestamp alone never invents a cadence. /// Upstream `GrokCreditsProxyFetcher.windowMinutes` keeps the length of a /// period that has already ended. -fn window_minutes( +pub(super) fn window_minutes( start: Option<&str>, end: Option>, now: DateTime, diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index 7c3f8db6239..923a6a879ce 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -5,6 +5,7 @@ //! gRPC-web endpoint. Browser cookies use the gRPC-web endpoint only. pub mod accounts; +mod acp; mod billing; mod credits_proxy; pub mod local_sessions; @@ -36,6 +37,8 @@ const REMAINING_RESETS_ENDPOINT: &str = const CLI_SETTINGS_ENDPOINT: &str = "https://cli-chat-proxy.grok.com/v1/settings"; const RESET_CREDITS_TIMEOUT: Duration = Duration::from_secs(2); const RESET_CREDITS_JOIN_GRACE: Duration = Duration::from_millis(250); +const GROK_CLI_PROGRAM: &str = "grok"; +const GROK_CLI_ACP_ARGS: [&str; 2] = ["agent", "stdio"]; /// Reset-window pace (upstream `GrokProviderDescriptor.pace`, v0.70.0): weekly /// windows inside their length, see [`reset_window_pace_matches`]. @@ -48,6 +51,33 @@ pub struct GrokProvider { client: Client, billing_endpoint: String, credits_proxy_endpoint: String, + cli_billing: CliBillingRunner, +} + +/// Where the ACP billing fallback gets its answer. `Installed` runs the Grok +/// CLI; tests substitute a fake to pin the fallback ordering. +#[derive(Clone)] +enum CliBillingRunner { + Installed, + #[cfg(test)] + Fake(std::sync::Arc Result + Send + Sync>), +} + +impl CliBillingRunner { + async fn fetch(&self) -> Result { + match self { + Self::Installed => { + acp::fetch_billing( + GROK_CLI_PROGRAM, + &GROK_CLI_ACP_ARGS, + acp::AcpTimeouts::default(), + ) + .await + } + #[cfg(test)] + Self::Fake(run) => run(), + } + } } impl GrokProvider { @@ -59,6 +89,7 @@ impl GrokProvider { .unwrap_or_else(|_| Client::new()), billing_endpoint: BILLING_ENDPOINT.to_string(), credits_proxy_endpoint: CREDITS_PROXY_ENDPOINT.to_string(), + cli_billing: CliBillingRunner::Installed, } } @@ -88,6 +119,15 @@ impl GrokProvider { self } + #[cfg(test)] + fn with_cli_billing_for_tests( + mut self, + run: impl Fn() -> Result + Send + Sync + 'static, + ) -> Self { + self.cli_billing = CliBillingRunner::Fake(std::sync::Arc::new(run)); + self + } + async fn load_credentials(kind: GrokAuthKind) -> Result { if kind == GrokAuthKind::OAuth { let text = accounts::AccountManager::new() @@ -124,7 +164,7 @@ impl GrokProvider { }; let credentials = GrokCredentials::parse_for_kind(text, kind)?; let result = self - .fetch_with_auth(&credentials, kind, &FetchContext::default()) + .fetch_with_auth(&credentials, kind, &saved_account_usage_context()) .await?; Ok(account_usage_from_result(&result)) } @@ -149,7 +189,10 @@ impl GrokProvider { billing, subscription_tier, prepaid_balance_usd, - } = match self.fetch_bearer_billing(credentials).await { + } = match self + .fetch_bearer_or_cli_billing(credentials, kind, ctx) + .await + { Ok(bearer) => bearer, Err(error) => { reset_lookup.abort(); @@ -186,6 +229,38 @@ impl GrokProvider { }) } + /// Current billing for a credential. The bearer path (credits proxy, then + /// grok.com gRPC-web) runs first. Only when it fails, for the ambient Grok + /// CLI login, is the installed CLI asked over ACP (upstream #4595). An + /// isolated fetch never reaches the ambient login, so it gets the bearer + /// error unchanged. + async fn fetch_bearer_or_cli_billing( + &self, + credentials: &GrokCredentials, + kind: GrokAuthKind, + ctx: &FetchContext, + ) -> Result { + let error = match self.fetch_bearer_billing(credentials).await { + Ok(bearer) => return Ok(bearer), + Err(error) => error, + }; + if kind != GrokAuthKind::Cli || ctx.token_account_isolated { + return Err(error); + } + match self.cli_billing.fetch().await { + Ok(bearer) => { + tracing::debug!("Grok billing answered by the CLI ACP fallback"); + Ok(bearer) + } + Err(acp_error) => { + // The original error stays the answer, so auto mode still + // moves to its next step exactly as before. + tracing::debug!("Grok CLI ACP billing fallback failed: {acp_error}"); + Err(error) + } + } + } + async fn fetch_cli_subscription_tier(&self, credentials: &GrokCredentials) -> Option { let response = self .client @@ -693,6 +768,16 @@ fn grok_steps_for_context( steps } +/// A saved account row is one principal. Its auth text may not be the account +/// switched into `~/.grok`, so the fetch is marked isolated: it must never +/// reach the ambient CLI login, including over ACP. +fn saved_account_usage_context() -> FetchContext { + FetchContext { + token_account_isolated: true, + ..FetchContext::default() + } +} + fn isolated_account_missing_credential() -> ProviderError { ProviderError::Other( "This Grok account has no saved token or cookie. Add one in Settings.".to_string(), @@ -996,3 +1081,7 @@ mod credits_proxy_tests; #[cfg(test)] #[path = "product_usage_tests.rs"] mod product_usage_tests; + +#[cfg(test)] +#[path = "cli_fallback_tests.rs"] +mod cli_fallback_tests; From e682f82b96d63b3904c28be1338a8ca3d7bd1864 Mon Sep 17 00:00:00 2001 From: Factory Agent Date: Wed, 21 Oct 2026 12:25:43 +0700 Subject: [PATCH 2/3] Document Grok ACP billing fallback Describe the Grok billing source order, the ACP fallback limits, its timeouts and the team-principal diagnostic that is not ported. --- docs/PROVIDERS.md | 40 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 40 insertions(+) diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index 829d1424e39..ed267b59fab 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -824,6 +824,46 @@ Origin: https://app.cosmicjs.com - 401, 403, and 429 report their status; response bodies are never echoed in errors. +## Grok + +Grok reads its usage from the Grok CLI login in `~/.grok/auth.json`, a saved +account row, a token account, or a browser session. A bearer login tries these +sources in order and stops at the first one that answers: + +1. Credits proxy: `GET https://cli-chat-proxy.grok.com/v1/billing?format=credits` + with the login's bearer token. If the answer has no quota percent, grok.com + gRPC-web may supply one within 6 seconds. +2. grok.com gRPC-web billing (`POST .../GetGrokCreditsConfig`), used when the + proxy fails. +3. Grok CLI ACP billing, a fallback that runs only when 1 and 2 both fail. + +The ACP fallback (upstream #4595) runs only for the ambient Grok CLI login. It +never runs for an isolated token account, a saved account row, or an OAuth +principal, so it can never read a different account's usage: + +- The installed `grok agent stdio` is started with no console window and its + stderr discarded. Requests are newline-delimited JSON-RPC 2.0 on stdin and + stdout. +- `initialize` is sent first with `protocolVersion` `"1"` and filesystem and + terminal capabilities off, with a 4-second budget. Then `_x.ai/billing` is + sent with empty params, with a 3-second budget. The method name is written + literally, without a `\/` escape. +- The child is killed after every call, including on timeout. It is never left + running. +- A current result has the nested `config` credits shape, read by the same + parser as the credits proxy. A top-level `subscription_tier` names the plan + and overrides `config.subscriptionTier`. The older top-level shape with + `monthlyLimit` and `usage.totalUsed` is still read. +- A prepaid balance counts as an answer without a quota percent. A result with + neither a percent nor a balance is unavailable. +- A JSON-RPC `-32601` error means this CLI does not provide billing over ACP. + An error that says authentication is required, or mentions `grok login`, is + an auth failure. In both cases the bearer path's original error is shown, so + auto mode continues to its next source as before. +- Not ported: upstream's team-principal diagnostic, which turns a `-32601` + answer for a team login into an identity-only row with an unsupported-team + usage message. On Windows that case shows the original error. + ## Window lengths that match the Mac These providers set window lengths so their rows get the same pace as in the Mac app, and the tray tint can score them. Where the stored length differs from the Mac's, the row says so: From 8529cd4b23d253a9070fc2abeff43145199dfd4f Mon Sep 17 00:00:00 2001 From: Factory Agent Date: Wed, 21 Oct 2026 12:48:26 +0700 Subject: [PATCH 3/3] Ask Grok CLI over ACP when its login is missing or expired --- docs/PROVIDERS.md | 8 +- rust/src/providers/grok/acp.rs | 5 +- rust/src/providers/grok/cli_fallback_tests.rs | 169 ++++++++++++++++++ rust/src/providers/grok/mod.rs | 86 ++++++++- 4 files changed, 258 insertions(+), 10 deletions(-) diff --git a/docs/PROVIDERS.md b/docs/PROVIDERS.md index ed267b59fab..e4b040d30ac 100644 --- a/docs/PROVIDERS.md +++ b/docs/PROVIDERS.md @@ -838,8 +838,12 @@ sources in order and stops at the first one that answers: 3. Grok CLI ACP billing, a fallback that runs only when 1 and 2 both fail. The ACP fallback (upstream #4595) runs only for the ambient Grok CLI login. It -never runs for an isolated token account, a saved account row, or an OAuth -principal, so it can never read a different account's usage: +also runs when that login is missing (no `~/.grok/auth.json` or no Cli entry) or +expired, since the installed CLI may hold its own valid session. In that case +there is no CodexBar token, so the answer has no account identity and no +reset-credit lookup. If ACP also fails, the original login error is shown, not +an ACP error. It never runs for an isolated token account, a saved account row, +or an OAuth principal, so it can never read a different account's usage: - The installed `grok agent stdio` is started with no console window and its stderr discarded. Requests are newline-delimited JSON-RPC 2.0 on stdin and diff --git a/rust/src/providers/grok/acp.rs b/rust/src/providers/grok/acp.rs index 3fb3095d0d1..56202e6d750 100644 --- a/rust/src/providers/grok/acp.rs +++ b/rust/src/providers/grok/acp.rs @@ -5,8 +5,9 @@ //! answers from its own login, so no token from CodexBar is sent. //! //! This is a fallback only. `mod.rs` calls it after the bearer billing path has -//! failed, and only for the ambient Grok CLI login. Isolated token accounts and -//! saved account rows never reach it. +//! failed, or when the ambient Grok CLI login is missing or expired, and only +//! for the ambient Grok CLI login. Isolated token accounts and saved account +//! rows never reach it. use std::io; use std::process::Stdio; diff --git a/rust/src/providers/grok/cli_fallback_tests.rs b/rust/src/providers/grok/cli_fallback_tests.rs index 56823a6306a..52cc2243229 100644 --- a/rust/src/providers/grok/cli_fallback_tests.rs +++ b/rust/src/providers/grok/cli_fallback_tests.rs @@ -221,3 +221,172 @@ async fn saved_account_row_failure_never_reaches_the_ambient_cli() { assert_eq!(calls.load(Ordering::SeqCst), 0); assert!(matches!(error, ProviderError::Other(_)), "{error}"); } + +fn missing_login() -> ProviderError { + ProviderError::NotInstalled("Grok auth.json not found. Run `grok login`.".to_string()) +} + +fn expired_login() -> ProviderError { + ProviderError::AuthRequired +} + +#[test] +fn expired_cli_login_in_auth_json_is_auth_required() { + // The real expired-login error must qualify for the ACP fallback. + let text = r#"{"https://accounts.x.ai/sign-in": {"key": "cli-token", "auth_mode": "session", "expires_at": "2000-01-01T00:00:00Z"}}"#; + let error = GrokCredentials::parse_for_kind(text, GrokAuthKind::Cli).unwrap_err(); + + assert!(matches!(error, ProviderError::AuthRequired), "{error}"); + assert!(is_missing_or_expired_login(&error)); +} + +#[tokio::test] +async fn missing_cli_login_is_answered_by_the_cli() { + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = GrokProvider::new().with_cli_billing_for_tests(run); + + let result = provider + .fetch_cli_acp_without_login(missing_login(), &FetchContext::default()) + .await + .unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(result.usage.primary.used_percent, 64.0); + assert_eq!(result.source_label, "grok-cli"); + assert_eq!(result.usage.account_email, None); + assert_eq!( + result.usage.login_method.as_deref(), + Some("SuperGrok Heavy") + ); +} + +#[tokio::test] +async fn expired_cli_login_is_answered_by_the_cli() { + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = GrokProvider::new().with_cli_billing_for_tests(run); + + let result = provider + .fetch_cli_acp_without_login(expired_login(), &FetchContext::default()) + .await + .unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert_eq!(result.usage.primary.used_percent, 64.0); +} + +#[tokio::test] +async fn auto_mode_uses_the_cli_when_the_login_is_missing_or_expired() { + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = GrokProvider::new().with_cli_billing_for_tests(run); + + let missing = provider + .try_ambient_loaded( + GrokAuthKind::Cli, + Err(missing_login()), + &FetchContext::default(), + ) + .await + .expect("auto step answers") + .unwrap(); + let expired = provider + .try_ambient_loaded( + GrokAuthKind::Cli, + Err(expired_login()), + &FetchContext::default(), + ) + .await + .expect("auto step answers") + .unwrap(); + + assert_eq!(calls.load(Ordering::SeqCst), 2); + assert_eq!(missing.usage.primary.used_percent, 64.0); + assert_eq!(expired.usage.primary.used_percent, 64.0); +} + +#[tokio::test] +async fn failing_cli_acp_keeps_the_original_missing_login_error() { + // The CLI's own auth error must not replace the friendly login error. + let (calls, run) = counting_cli(|| Err(ProviderError::AuthRequired)); + let provider = GrokProvider::new().with_cli_billing_for_tests(run); + + let error = provider + .fetch_cli_acp_without_login(missing_login(), &FetchContext::default()) + .await + .unwrap_err(); + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!( + matches!(&error, ProviderError::NotInstalled(message) if message.contains("grok login")), + "{error}" + ); +} + +#[tokio::test] +async fn failing_cli_acp_in_auto_mode_moves_to_the_next_step() { + let (calls, run) = counting_cli(|| Err(ProviderError::AuthRequired)); + let provider = GrokProvider::new().with_cli_billing_for_tests(run); + + let step = provider + .try_ambient_loaded( + GrokAuthKind::Cli, + Err(missing_login()), + &FetchContext::default(), + ) + .await; + + assert_eq!(calls.load(Ordering::SeqCst), 1); + assert!(step.is_none()); +} + +#[tokio::test] +async fn isolated_token_account_never_reaches_the_cli_without_a_login() { + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = GrokProvider::new().with_cli_billing_for_tests(run); + let ctx = FetchContext { + token_account_isolated: true, + ..FetchContext::default() + }; + + let error = provider + .fetch_cli_acp_without_login(missing_login(), &ctx) + .await + .unwrap_err(); + + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert!(matches!(error, ProviderError::NotInstalled(_)), "{error}"); +} + +#[tokio::test] +async fn malformed_cli_login_never_reaches_the_cli() { + // A parse failure is a broken file, not a missing or expired login. + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = GrokProvider::new().with_cli_billing_for_tests(run); + + let error = provider + .fetch_cli_acp_without_login( + ProviderError::Parse("bad json".to_string()), + &FetchContext::default(), + ) + .await + .unwrap_err(); + + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert!(matches!(error, ProviderError::Parse(_)), "{error}"); +} + +#[tokio::test] +async fn oauth_login_missing_never_reaches_the_cli() { + let (calls, run) = counting_cli(|| Ok(cli_answer())); + let provider = GrokProvider::new().with_cli_billing_for_tests(run); + + let step = provider + .try_ambient_loaded( + GrokAuthKind::OAuth, + Err(missing_login()), + &FetchContext::default(), + ) + .await; + + assert_eq!(calls.load(Ordering::SeqCst), 0); + assert!(step.is_none()); +} diff --git a/rust/src/providers/grok/mod.rs b/rust/src/providers/grok/mod.rs index 923a6a879ce..26f8e7afa62 100644 --- a/rust/src/providers/grok/mod.rs +++ b/rust/src/providers/grok/mod.rs @@ -261,6 +261,31 @@ impl GrokProvider { } } + /// The ambient Grok CLI login is missing or expired, so no credential can + /// be sent. The installed CLI may still hold its own valid session, so it + /// is asked over ACP as a fallback. Isolated fetches never reach the + /// ambient login. When ACP also fails, the original login error is + /// returned, not an ACP error. + async fn fetch_cli_acp_without_login( + &self, + login_error: ProviderError, + ctx: &FetchContext, + ) -> Result { + if ctx.token_account_isolated || !is_missing_or_expired_login(&login_error) { + return Err(login_error); + } + match self.cli_billing.fetch().await { + Ok(bearer) => { + tracing::debug!("Grok billing answered by the CLI ACP fallback without a login"); + Ok(result_from_acp_billing(bearer)) + } + Err(acp_error) => { + tracing::debug!("Grok CLI ACP billing fallback failed: {acp_error}"); + Err(login_error) + } + } + } + async fn fetch_cli_subscription_tier(&self, credentials: &GrokCredentials) -> Option { let response = self .client @@ -414,7 +439,29 @@ impl GrokProvider { kind: GrokAuthKind, ctx: &FetchContext, ) -> Option> { - let credentials = Self::load_credentials(kind).await.ok()?; + let loaded = Self::load_credentials(kind).await; + self.try_ambient_loaded(kind, loaded, ctx).await + } + + /// Auto-mode step for an ambient login that has already been loaded (or + /// failed to load). A missing or expired Cli login may still be answered by + /// the CLI over ACP; if that fails, the step yields `None` as before. + async fn try_ambient_loaded( + &self, + kind: GrokAuthKind, + loaded: Result, + ctx: &FetchContext, + ) -> Option> { + let credentials = match loaded { + Ok(credentials) => credentials, + Err(error) if kind == GrokAuthKind::Cli => { + return match self.fetch_cli_acp_without_login(error, ctx).await { + Ok(result) => Some(Ok(result)), + Err(_) => None, + }; + } + Err(_) => return None, + }; match self.fetch_with_auth(&credentials, kind, ctx).await { Ok(result) => Some(Ok(result)), Err(ProviderError::AuthRequired) => None, @@ -681,11 +728,13 @@ impl Provider for GrokProvider { SourceMode::Cli | SourceMode::OAuth if ctx.token_account_isolated => { self.fetch_isolated_account(ctx).await } - SourceMode::Cli => { - let credentials = Self::load_credentials(GrokAuthKind::Cli).await?; - self.fetch_with_auth(&credentials, GrokAuthKind::Cli, ctx) - .await - } + SourceMode::Cli => match Self::load_credentials(GrokAuthKind::Cli).await { + Ok(credentials) => { + self.fetch_with_auth(&credentials, GrokAuthKind::Cli, ctx) + .await + } + Err(error) => self.fetch_cli_acp_without_login(error, ctx).await, + }, SourceMode::OAuth => { // Prefer the switched ~/.grok/auth.json over a leftover token // account so Weekly/notifications follow Grok account Switch. @@ -1043,6 +1092,31 @@ fn result_from_billing( ) } +/// Result for a CLI login-less ACP answer. The CLI answered from its own +/// login, so no account identity is known and no reset-credit lookup runs +/// (that lookup needs the token CodexBar does not have). +fn result_from_acp_billing(bearer: BearerBilling) -> ProviderFetchResult { + let BearerBilling { + billing, + subscription_tier, + prepaid_balance_usd, + } = bearer; + attach_prepaid_balance( + result_from_billing(billing, "grok-cli", None, None, subscription_tier), + prepaid_balance_usd, + ) +} + +/// A Cli login that is absent (no auth file or no Cli entry) or expired. +/// Parse errors are a malformed file, not a missing login, so they do not +/// qualify for the ACP fallback. +fn is_missing_or_expired_login(err: &ProviderError) -> bool { + matches!( + err, + ProviderError::AuthRequired | ProviderError::NotInstalled(_) + ) +} + /// Whether a cookie-path error should invalidate the cached browser session. fn is_cookie_authentication_failure(err: &ProviderError) -> bool { matches!(err, ProviderError::AuthRequired)