diff --git a/crates/openab-core/src/acp/connection.rs b/crates/openab-core/src/acp/connection.rs index fbe225df8..f40ceb486 100644 --- a/crates/openab-core/src/acp/connection.rs +++ b/crates/openab-core/src/acp/connection.rs @@ -1,5 +1,6 @@ use crate::acp::protocol::{ - parse_config_options, ConfigOption, JsonRpcMessage, JsonRpcRequest, JsonRpcResponse, + parse_config_options, parse_usage_report, ConfigOption, JsonRpcMessage, JsonRpcRequest, + JsonRpcResponse, UsageReport, }; use anyhow::{anyhow, Result}; use serde_json::{json, Value}; @@ -179,6 +180,9 @@ pub struct AcpConnection { notify_tx: Arc>>>, pub acp_session_id: Option, pub supports_load_session: bool, + /// Agent name from `initialize` (`agentInfo.name`), e.g. "Kiro CLI Agent". + /// Used to gate agent-specific extension methods. + pub agent_name: String, pub config_options: Vec, pub last_active: Instant, pub activity: Arc, @@ -474,6 +478,7 @@ impl AcpConnection { notify_tx, acp_session_id: None, supports_load_session: false, + agent_name: String::new(), config_options: Vec::new(), last_active: Instant::now(), activity, @@ -543,6 +548,7 @@ impl AcpConnection { .and_then(|a| a.get("name")) .and_then(|n| n.as_str()) .unwrap_or("unknown"); + self.agent_name = agent_name.to_string(); self.supports_load_session = result .and_then(|r| r.get("agentCapabilities")) .and_then(|c| c.get("loadSession")) @@ -638,6 +644,50 @@ impl AcpConnection { Ok(self.config_options.clone()) } + /// Query account-level usage/billing via kiro-cli's + /// `_kiro.dev/commands/execute` extension (the `/usage` slash command). + /// + /// This is a Kiro-specific extension, not part of the ACP spec, and the + /// request shape is strict: a malformed `command` value is a + /// deserialization error that kills the whole ACP connection (no JSON-RPC + /// error is returned). We therefore gate on the agent name advertised in + /// `initialize` and never retry on failure. + pub async fn get_usage(&mut self) -> Result { + if !self.agent_name.to_ascii_lowercase().contains("kiro") { + return Err(anyhow!( + "usage query is not supported by this backend ({})", + if self.agent_name.is_empty() { + "unknown agent" + } else { + &self.agent_name + } + )); + } + let session_id = self + .acp_session_id + .as_ref() + .ok_or_else(|| anyhow!("no session"))? + .clone(); + + let resp = self + .send_request( + "_kiro.dev/commands/execute", + Some(json!({ + "sessionId": session_id, + // Adjacently-tagged TuiCommand enum: tag = "command", content = "args". + "command": {"command": "usage", "args": {}}, + })), + ) + .await?; + + let result = resp + .result + .as_ref() + .ok_or_else(|| anyhow!("empty usage response"))?; + parse_usage_report(result) + .ok_or_else(|| anyhow!("could not parse usage response from agent")) + } + /// Send a prompt with content blocks (text and/or images) and return a receiver /// for streaming notifications. The final message on the channel will have id set /// (the prompt response). diff --git a/crates/openab-core/src/acp/pool.rs b/crates/openab-core/src/acp/pool.rs index 16497b324..394d9f260 100644 --- a/crates/openab-core/src/acp/pool.rs +++ b/crates/openab-core/src/acp/pool.rs @@ -568,6 +568,22 @@ impl SessionPool { conn.set_config_option(config_id, value).await } + /// Query account-level usage/billing from the backend agent for a session + /// (kiro-cli extension). Fails when there is no active session for the + /// thread or the backend does not support usage queries. + pub async fn get_usage(&self, thread_id: &str) -> Result { + let conn = { + let state = self.state.read().await; + state + .active + .get(thread_id) + .cloned() + .ok_or_else(|| anyhow!("no connection for thread {thread_id}"))? + }; + let mut conn = conn.lock().await; + conn.get_usage().await + } + /// Cancel the current in-flight operation for a session. /// Uses pre-stored cancel handles to avoid locking the connection (which is held during streaming). pub async fn cancel_session(&self, thread_id: &str) -> Result<()> { diff --git a/crates/openab-core/src/acp/protocol.rs b/crates/openab-core/src/acp/protocol.rs index 0f89b8f59..4353642b2 100644 --- a/crates/openab-core/src/acp/protocol.rs +++ b/crates/openab-core/src/acp/protocol.rs @@ -241,6 +241,99 @@ pub fn parse_turn_result(result: &Value) -> TurnResult { } } +// --- Account usage report (kiro-cli `_kiro.dev/commands/execute` extension) --- + +/// One resource-type usage breakdown (e.g. credits) from a usage report. +#[derive(Debug, Clone, PartialEq)] +pub struct UsageBreakdown { + pub display_name: String, + pub used: f64, + /// Plan allowance for this resource. `None` when the account has no + /// per-user cap (`hasLimit: false`, e.g. pooled enterprise credits). + pub limit: Option, + pub percentage: Option, + /// Accrued overage charges for this cycle, if any. + pub overage_charges: Option, + pub currency: Option, +} + +/// Account-level usage/billing report returned by kiro-cli's `/usage` command +/// when executed over ACP via `_kiro.dev/commands/execute`. +#[derive(Debug, Clone, PartialEq)] +pub struct UsageReport { + pub plan_name: String, + /// Billing cycle reset date (e.g. "2026-08-01"), if reported. + pub billing_cycle_reset: Option, + pub breakdowns: Vec, +} + +/// Parse a usage report from the result of a `_kiro.dev/commands/execute` +/// request for the `usage` command. Expected shape (kiro-cli 2.12.x): +/// +/// ```json +/// {"success": true, "message": "...", "data": { +/// "planName": "KIRO POWER", "billingCycleReset": "2026-08-01", +/// "usageBreakdowns": [{"displayName": "Credits", "used": 128.5, +/// "limit": 10000.0, "percentage": 1, "hasLimit": true, +/// "overageCharges": 0.0, "currency": "USD"}]}} +/// ``` +/// +/// Returns `None` when `success` is not true or the data shape is missing — +/// callers should treat that as "usage not supported by this agent". +pub fn parse_usage_report(result: &Value) -> Option { + if !result.get("success").and_then(|v| v.as_bool()).unwrap_or(false) { + return None; + } + let data = result.get("data")?; + let plan_name = data.get("planName")?.as_str()?.to_string(); + let billing_cycle_reset = data + .get("billingCycleReset") + .and_then(|v| v.as_str()) + .map(String::from); + + let breakdowns = data + .get("usageBreakdowns") + .and_then(|v| v.as_array()) + .map(|items| { + items + .iter() + .filter_map(|b| { + // `hasLimit: false` is an explicit sentinel (pooled/no-cap + // enterprise accounts). Some kiro-cli versions omit the + // field entirely while still returning a numeric `limit` + // (observed live: overage state on kiro-cli in the pr1392 + // image) — treat a missing `hasLimit` as "has a limit if a + // numeric limit is present" instead of hiding the cap. + let has_limit = b + .get("hasLimit") + .and_then(|v| v.as_bool()) + .unwrap_or_else(|| { + b.get("limit").and_then(|v| v.as_f64()).is_some() + }); + Some(UsageBreakdown { + display_name: b.get("displayName")?.as_str()?.to_string(), + used: b.get("used")?.as_f64()?, + limit: if has_limit { + b.get("limit").and_then(|v| v.as_f64()) + } else { + None + }, + percentage: b.get("percentage").and_then(|v| v.as_u64()), + overage_charges: b.get("overageCharges").and_then(|v| v.as_f64()), + currency: b.get("currency").and_then(|v| v.as_str()).map(String::from), + }) + }) + .collect() + }) + .unwrap_or_default(); + + Some(UsageReport { + plan_name, + billing_cycle_reset, + breakdowns, + }) +} + // --- ACP notification classification --- #[derive(Debug)] @@ -495,4 +588,131 @@ mod tests { let tr = parse_turn_result(&result); assert!(!tr.is_silent_failure()); } + + #[test] + fn parse_usage_report_full() { + let result = json!({ + "success": true, + "message": "Plan: KIRO POWER | 1 usage breakdowns", + "data": { + "planName": "KIRO POWER", + "billingCycleReset": "2026-08-01", + "overagesEnabled": true, + "isEnterprise": false, + "usageBreakdowns": [{ + "resourceType": "CREDIT", + "displayName": "Credits", + "used": 12781.64, + "limit": 10000.0, + "percentage": 127, + "currentOverages": 2781.64, + "overageRate": 0.04, + "overageCharges": 111.27, + "currency": "USD", + "hasLimit": true + }], + "bonusCredits": [], + "addOnCredits": [], + "overageCapable": true + } + }); + let report = parse_usage_report(&result).expect("should parse"); + assert_eq!(report.plan_name, "KIRO POWER"); + assert_eq!(report.billing_cycle_reset.as_deref(), Some("2026-08-01")); + assert_eq!(report.breakdowns.len(), 1); + let b = &report.breakdowns[0]; + assert_eq!(b.display_name, "Credits"); + assert_eq!(b.used, 12781.64); + assert_eq!(b.limit, Some(10000.0)); + assert_eq!(b.percentage, Some(127)); + assert_eq!(b.overage_charges, Some(111.27)); + assert_eq!(b.currency.as_deref(), Some("USD")); + } + + #[test] + fn parse_usage_report_no_limit_hides_cap() { + // Pooled enterprise credits: hasLimit=false means the backend's + // sentinel limit value must not be surfaced. + let result = json!({ + "success": true, + "data": { + "planName": "ENTERPRISE", + "usageBreakdowns": [{ + "displayName": "Credits", + "used": 320.0, + "limit": 999999.0, + "hasLimit": false + }] + } + }); + let report = parse_usage_report(&result).expect("should parse"); + assert_eq!(report.breakdowns[0].limit, None); + assert_eq!(report.billing_cycle_reset, None); + } + + #[test] + fn parse_usage_report_missing_has_limit_keeps_numeric_limit() { + // Regression: exact payload captured live from B0 (kiro-cli in the + // pr1392 image, 2026-07-13). This kiro-cli version omits `hasLimit` + // entirely while returning a real numeric `limit` — the old + // `unwrap_or(false)` default hid the cap and progress bar even though + // the account was 130% over its 10000-credit limit. + let result = json!({ + "success": true, + "message": "Plan: KIRO POWER | 1 usage breakdowns", + "data": { + "planName": "KIRO POWER", + "billingCycleReset": "2026-08-01", + "overagesEnabled": true, + "isEnterprise": false, + "usageBreakdowns": [{ + "resourceType": "CREDIT", + "displayName": "Credits", + "used": 13080.52, + "limit": 10000.0, + "percentage": 130, + "currentOverages": 3080.52, + "overageRate": 0.04, + "overageCharges": 123.220870825752, + "currency": "USD" + }], + "bonusCredits": [] + } + }); + let report = parse_usage_report(&result).expect("should parse"); + let b = &report.breakdowns[0]; + assert_eq!(b.limit, Some(10000.0)); + assert_eq!(b.percentage, Some(130)); + } + + #[test] + fn parse_usage_report_missing_has_limit_and_missing_limit_hides_cap() { + // No `hasLimit` and no numeric `limit` → still consumption-only. + let result = json!({ + "success": true, + "data": { + "planName": "POOLED", + "usageBreakdowns": [{"displayName": "Credits", "used": 320.0}] + } + }); + let report = parse_usage_report(&result).expect("should parse"); + assert_eq!(report.breakdowns[0].limit, None); + } + + #[test] + fn parse_usage_report_failure_returns_none() { + assert!(parse_usage_report(&json!({"success": false})).is_none()); + assert!(parse_usage_report(&json!({"success": true})).is_none()); // no data + assert!(parse_usage_report(&json!({})).is_none()); + } + + #[test] + fn parse_usage_report_empty_breakdowns() { + let result = json!({ + "success": true, + "data": {"planName": "FREE", "usageBreakdowns": []} + }); + let report = parse_usage_report(&result).expect("should parse"); + assert!(report.breakdowns.is_empty()); + } } diff --git a/crates/openab-core/src/discord.rs b/crates/openab-core/src/discord.rs index be0b0a44e..0489ce83b 100644 --- a/crates/openab-core/src/discord.rs +++ b/crates/openab-core/src/discord.rs @@ -1,4 +1,4 @@ -use crate::acp::protocol::ConfigOption; +use crate::acp::protocol::{ConfigOption, UsageReport}; use crate::acp::ContentBlock; use crate::adapter::{AdapterRouter, ChannelRef, ChatAdapter, MessageRef, SenderContext}; use crate::bot_turns::{BotTurnTracker, TurnAction, TurnSeverity, BOT_TURN_LIMIT_WARNING_PREFIX}; @@ -11,9 +11,9 @@ use crate::trust::l3_gate_applies; use async_trait::async_trait; use serenity::builder::{ CreateActionRow, CreateAttachment, CreateButton, CreateCommand, CreateCommandOption, - CreateInteractionResponse, CreateInteractionResponseFollowup, CreateInteractionResponseMessage, - CreateSelectMenu, CreateSelectMenuKind, CreateSelectMenuOption, CreateThread, EditChannel, - EditMessage, GetMessages, + CreateEmbed, CreateEmbedFooter, CreateInteractionResponse, CreateInteractionResponseFollowup, + CreateInteractionResponseMessage, CreateSelectMenu, CreateSelectMenuKind, + CreateSelectMenuOption, CreateThread, EditChannel, EditMessage, GetMessages, }; use serenity::http::Http; use serenity::model::application::ButtonStyle; @@ -1398,6 +1398,8 @@ impl EventHandler for Handler { ).required(true)), CreateCommand::new("auth") .description("Authenticate the backend agent (device flow)"), + CreateCommand::new("usage") + .description("Show backend account usage and billing information"), CreateCommand::new("export-thread") .description("Download this thread as a text file") .add_option(CreateCommandOption::new( @@ -1489,6 +1491,9 @@ impl EventHandler for Handler { Interaction::Command(cmd) if cmd.data.name == "auth" => { self.handle_auth_command(&ctx, &cmd).await; } + Interaction::Command(cmd) if cmd.data.name == "usage" => { + self.handle_usage_command(&ctx, &cmd).await; + } Interaction::Component(comp) if comp.data.custom_id.starts_with("acp_config_") => { self.handle_config_select(&ctx, &comp).await; } @@ -1655,6 +1660,51 @@ impl Handler { } } + async fn handle_usage_command( + &self, + ctx: &Context, + cmd: &serenity::model::application::CommandInteraction, + ) { + let thread_key = format!("discord:{}", cmd.channel_id.get()); + + if !self.router.pool().has_active_session(&thread_key).await { + let response = CreateInteractionResponse::Message( + CreateInteractionResponseMessage::new() + .content("⚠️ No active session. Start a conversation first by @mentioning the bot.") + .ephemeral(true), + ); + if let Err(e) = cmd.create_response(&ctx.http, response).await { + tracing::error!(error = %e, "failed to respond to /usage command"); + } + return; + } + + // The ACP round-trip can exceed Discord's 3-second interaction + // deadline — acknowledge with a deferred ephemeral response first. + let defer = + CreateInteractionResponse::Defer(CreateInteractionResponseMessage::new().ephemeral(true)); + if let Err(e) = cmd.create_response(&ctx.http, defer).await { + tracing::error!(error = %e, "failed to defer /usage response"); + return; + } + + let followup = match self.router.pool().get_usage(&thread_key).await { + Ok(report) => { + let (content, embed) = build_usage_reply(&report); + CreateInteractionResponseFollowup::new() + .content(content) + .embed(embed) + .ephemeral(true) + } + Err(e) => CreateInteractionResponseFollowup::new() + .content(format!("⚠️ {e}")) + .ephemeral(true), + }; + if let Err(e) = cmd.create_followup(&ctx.http, followup).await { + tracing::error!(error = %e, "failed to send /usage followup"); + } + } + async fn handle_cancel_command( &self, ctx: &Context, @@ -2473,6 +2523,75 @@ impl Handler { // --- Discord-specific helpers --- +/// Render the body lines of a usage report (everything except the plan title +/// and billing-cycle footer). Returns the text and whether any breakdown is +/// over its plan limit. +fn format_usage_body(report: &UsageReport) -> (String, bool) { + let mut lines: Vec = Vec::new(); + let mut over_limit = false; + + for b in &report.breakdowns { + match b.limit { + Some(limit) => { + let pct = b.percentage.unwrap_or_else(|| { + if limit > 0.0 { + (b.used / limit * 100.0).round() as u64 + } else { + 0 + } + }); + if pct > 100 { + over_limit = true; + } + // 10-slot progress bar, clamped at 100%. + let filled = (pct.min(100) as usize) / 10; + let bar: String = "█".repeat(filled) + &"░".repeat(10 - filled); + lines.push(format!( + "{}: {:.2} / {:.0} `{}` {}%{}", + b.display_name, + b.used, + limit, + bar, + pct, + if pct > 100 { " ⚠️" } else { "" } + )); + } + // No per-user cap (e.g. pooled enterprise credits). + None => lines.push(format!("{}: {:.2} used", b.display_name, b.used)), + } + if let Some(charges) = b.overage_charges { + if charges > 0.0 { + lines.push(format!( + "Overage charges: {:.2} {}", + charges, + b.currency.as_deref().unwrap_or("USD") + )); + } + } + } + + (lines.join("\n"), over_limit) +} + +/// Build the /usage reply as full-size message content plus a minimal +/// color-strip embed. Discord renders embed descriptions at a smaller font +/// than message content, so the report body lives in `content` (normal font) +/// while the embed carries only the at-a-glance color signal (green within +/// the plan limit, red when any breakdown is over) and the billing-cycle +/// footer. +fn build_usage_reply(report: &UsageReport) -> (String, CreateEmbed) { + let (body, over_limit) = format_usage_body(report); + let content = format!("📊 **Usage — {}**\n{}", report.plan_name, body); + let mut embed = + CreateEmbed::new().colour(if over_limit { 0xE74C3C } else { 0x2ECC71 }); + if let Some(reset) = &report.billing_cycle_reset { + embed = embed.footer(CreateEmbedFooter::new(format!( + "Billing cycle resets {reset}" + ))); + } + (content, embed) +} + fn discord_msg_ref(msg: &Message) -> MessageRef { MessageRef { channel: ChannelRef { @@ -3128,6 +3247,99 @@ mod tests { use super::*; use crate::bot_turns::{TurnResult, HARD_BOT_TURN_LIMIT, BOT_TURN_LIMIT_WARNING_PREFIX}; + // --- format_usage_report tests (/usage slash command) --- + + fn usage_breakdown() -> crate::acp::protocol::UsageBreakdown { + crate::acp::protocol::UsageBreakdown { + display_name: "Credits".into(), + used: 12781.64, + limit: Some(10000.0), + percentage: Some(127), + overage_charges: Some(111.27), + currency: Some("USD".into()), + } + } + + #[test] + fn format_usage_over_limit() { + let report = UsageReport { + plan_name: "KIRO POWER".into(), + billing_cycle_reset: Some("2026-08-01".into()), + breakdowns: vec![usage_breakdown()], + }; + let (body, over_limit) = format_usage_body(&report); + assert!(over_limit); + assert!(body.contains("12781.64 / 10000")); + assert!(body.contains("127%")); + assert!(body.contains("⚠️")); + assert!(body.contains("Overage charges: 111.27 USD")); + } + + /// The report body must ride in the message content (normal font size), + /// not the embed description (rendered smaller by Discord clients). The + /// embed only carries the color strip + billing-cycle footer. + #[test] + fn usage_reply_body_in_content_not_embed() { + let report = UsageReport { + plan_name: "KIRO POWER".into(), + billing_cycle_reset: Some("2026-08-01".into()), + breakdowns: vec![usage_breakdown()], + }; + let (content, embed) = build_usage_reply(&report); + assert!(content.starts_with("📊 **Usage — KIRO POWER**")); + assert!(content.contains("12781.64 / 10000")); + assert!(content.contains("Overage charges: 111.27 USD")); + let json = serde_json::to_value(&embed).expect("embed serializes"); + assert!(json.get("description").is_none(), "body must not be in embed"); + assert!(json.get("title").is_none(), "title must not be in embed"); + assert_eq!(json["color"], 0xE74C3C, "over limit → red strip"); + assert_eq!(json["footer"]["text"], "Billing cycle resets 2026-08-01"); + } + + #[test] + fn format_usage_no_limit_shows_consumption_only() { + let report = UsageReport { + plan_name: "ENTERPRISE".into(), + billing_cycle_reset: None, + breakdowns: vec![crate::acp::protocol::UsageBreakdown { + display_name: "Credits".into(), + used: 320.0, + limit: None, + percentage: None, + overage_charges: None, + currency: None, + }], + }; + let (body, over_limit) = format_usage_body(&report); + assert!(!over_limit); + assert!(body.contains("Credits: 320.00 used")); + assert!(!body.contains('/')); + assert!(!body.contains("Overage")); + } + + #[test] + fn format_usage_under_limit_no_warning() { + let report = UsageReport { + plan_name: "FREE".into(), + billing_cycle_reset: None, + breakdowns: vec![crate::acp::protocol::UsageBreakdown { + display_name: "Credits".into(), + used: 50.0, + limit: Some(100.0), + percentage: Some(50), + overage_charges: Some(0.0), + currency: Some("USD".into()), + }], + }; + let (body, over_limit) = format_usage_body(&report); + assert!(!over_limit); + assert!(body.contains("50%")); + assert!(!body.contains("⚠️")); + assert!(!body.contains("Overage")); + // 50% → 5 of 10 bar slots filled. + assert!(body.contains("█████░░░░░")); + } + // --- truncate_to_utf16_budget tests (#1185 /auth output relay) --- /// Body shorter than the budget is returned unchanged.