Skip to main content

ai_usagebar/kiro/
types.rs

1//! Wire schema for `AmazonCodeWhispererService.GetUsageLimits` — the exact
2//! call kiro-cli's own `/usage` slash command makes (confirmed by tracing
3//! `kiro-cli chat --no-interactive -vvv "/usage"`, which logs the operation
4//! name and the `management.<region>.kiro.dev` endpoint it resolves to for an
5//! IAM Identity Center account). Verified live against
6//! `https://codewhisperer.us-east-1.amazonaws.com/` with `x-amz-target:
7//! AmazonCodeWhispererService.GetUsageLimits` and the account's own cached
8//! bearer token:
9//!
10//! ```json
11//! {
12//!   "nextDateReset": 1785542400.0,
13//!   "subscriptionInfo": { "subscriptionTitle": "KIRO POWER" },
14//!   "usageBreakdownList": [{
15//!     "resourceType": "CREDIT",
16//!     "displayName": "Credit",
17//!     "currentUsageWithPrecision": 9943.38,
18//!     "usageLimitWithPrecision": 10000.0
19//!   }]
20//! }
21//! ```
22//!
23//! **Reverse-engineered, not documented** — CodeWhisperer/Q Developer has no
24//! public API reference (AWS's own docs say as much). Community reference
25//! implementations exist (`Finesssee/ProxyPilot`, `HsnSaboor/CLIProxyAPIPlus`)
26//! confirming the same request/response shape against `codewhisperer.*.amazonaws.com`.
27
28use chrono::{DateTime, Utc};
29use serde::Deserialize;
30
31use crate::error::{AppError, Result};
32use crate::usage::KiroSnapshot;
33
34#[derive(Debug, Clone, Deserialize)]
35#[serde(rename_all = "camelCase")]
36pub struct UsageLimitsResponse {
37    #[serde(default)]
38    pub subscription_info: Option<SubscriptionInfo>,
39    #[serde(default)]
40    pub usage_breakdown_list: Vec<UsageBreakdown>,
41    #[serde(default)]
42    pub next_date_reset: Option<f64>,
43}
44
45#[derive(Debug, Clone, Deserialize)]
46#[serde(rename_all = "camelCase")]
47pub struct SubscriptionInfo {
48    #[serde(default)]
49    pub subscription_title: Option<String>,
50}
51
52#[derive(Debug, Clone, Deserialize)]
53#[serde(rename_all = "camelCase")]
54pub struct UsageBreakdown {
55    #[serde(default)]
56    pub resource_type: Option<String>,
57    #[serde(default)]
58    pub current_usage_with_precision: Option<f64>,
59    #[serde(default)]
60    pub usage_limit_with_precision: Option<f64>,
61}
62
63/// The list can carry more than one resource bucket; the credit pool
64/// (`resourceType: "CREDIT"`) is what kiro-cli's own `/usage` renders, so it
65/// wins when present. A single-entry list with no `resourceType` (seen on
66/// some plans) is accepted as-is rather than rejected on a technicality.
67fn credit_breakdown(list: &[UsageBreakdown]) -> Result<&UsageBreakdown> {
68    if let Some(credit) = list
69        .iter()
70        .find(|b| b.resource_type.as_deref() == Some("CREDIT"))
71    {
72        return Ok(credit);
73    }
74    if let [only] = list
75        && only.resource_type.is_none()
76    {
77        return Ok(only);
78    }
79    if list.is_empty() {
80        Err(AppError::Schema(
81            "kiro: `usageBreakdownList` is empty".into(),
82        ))
83    } else {
84        Err(AppError::Schema(
85            "kiro: no unambiguous `CREDIT` usage bucket".into(),
86        ))
87    }
88}
89
90pub fn to_snapshot(resp: UsageLimitsResponse) -> Result<KiroSnapshot> {
91    let plan = resp
92        .subscription_info
93        .as_ref()
94        .and_then(|s| s.subscription_title.as_deref())
95        .filter(|s| !s.trim().is_empty())
96        .ok_or_else(|| {
97            AppError::Schema("kiro: missing `subscriptionInfo.subscriptionTitle`".into())
98        })?
99        .to_string();
100
101    let breakdown = credit_breakdown(&resp.usage_breakdown_list)?;
102
103    let used = finite(
104        "currentUsageWithPrecision",
105        breakdown.current_usage_with_precision,
106    )?;
107    let limit = finite(
108        "usageLimitWithPrecision",
109        breakdown.usage_limit_with_precision,
110    )?;
111
112    let reset_at = resp
113        .next_date_reset
114        .map(|secs| seconds_to_datetime("nextDateReset", secs))
115        .transpose()?;
116
117    Ok(KiroSnapshot {
118        plan,
119        used,
120        limit,
121        reset_at,
122    })
123}
124
125fn finite(field: &str, v: Option<f64>) -> Result<f64> {
126    let v = v.ok_or_else(|| AppError::Schema(format!("kiro: missing `{field}`")))?;
127    if !v.is_finite() || v < 0.0 {
128        return Err(AppError::Schema(format!(
129            "kiro: `{field}` is not a non-negative finite number ({v})"
130        )));
131    }
132    Ok(v)
133}
134
135fn seconds_to_datetime(field: &str, secs: f64) -> Result<DateTime<Utc>> {
136    if !secs.is_finite() || secs < 0.0 {
137        return Err(AppError::Schema(format!(
138            "kiro: `{field}` is not a valid Unix timestamp ({secs})"
139        )));
140    }
141    DateTime::from_timestamp(secs as i64, 0)
142        .ok_or_else(|| AppError::Schema(format!("kiro: `{field}` is out of range ({secs})")))
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    fn sample() -> UsageLimitsResponse {
150        serde_json::from_str(
151            r#"{
152                "daysUntilReset": 0,
153                "nextDateReset": 1785542400.0,
154                "subscriptionInfo": { "subscriptionTitle": "KIRO POWER" },
155                "usageBreakdownList": [{
156                    "resourceType": "CREDIT",
157                    "displayName": "Credit",
158                    "currentUsageWithPrecision": 9943.38,
159                    "usageLimitWithPrecision": 10000.0
160                }]
161            }"#,
162        )
163        .unwrap()
164    }
165
166    #[test]
167    fn parses_the_verified_live_shape() {
168        let snap = to_snapshot(sample()).unwrap();
169        assert_eq!(snap.plan, "KIRO POWER");
170        assert_eq!(snap.used, 9943.38);
171        assert_eq!(snap.limit, 10000.0);
172        assert_eq!(
173            snap.reset_at,
174            Some(DateTime::from_timestamp(1785542400, 0).unwrap())
175        );
176    }
177
178    #[test]
179    fn picks_the_credit_bucket_when_multiple_are_present() {
180        let mut resp = sample();
181        resp.usage_breakdown_list.insert(
182            0,
183            UsageBreakdown {
184                resource_type: Some("OTHER".into()),
185                current_usage_with_precision: Some(1.0),
186                usage_limit_with_precision: Some(2.0),
187            },
188        );
189        let snap = to_snapshot(resp).unwrap();
190        assert_eq!(snap.used, 9943.38);
191    }
192
193    #[test]
194    fn falls_back_to_the_first_entry_with_no_resource_type() {
195        let resp: UsageLimitsResponse = serde_json::from_str(
196            r#"{
197                "subscriptionInfo": { "subscriptionTitle": "KIRO POWER" },
198                "usageBreakdownList": [{
199                    "currentUsageWithPrecision": 5.0,
200                    "usageLimitWithPrecision": 10.0
201                }]
202            }"#,
203        )
204        .unwrap();
205        let snap = to_snapshot(resp).unwrap();
206        assert_eq!(snap.used, 5.0);
207        assert_eq!(snap.limit, 10.0);
208    }
209
210    #[test]
211    fn explicit_non_credit_single_bucket_is_schema_drift() {
212        let mut resp = sample();
213        resp.usage_breakdown_list[0].resource_type = Some("OTHER".into());
214        assert!(matches!(to_snapshot(resp), Err(AppError::Schema(_))));
215    }
216
217    #[test]
218    fn multiple_unknown_buckets_are_schema_drift() {
219        let mut resp = sample();
220        resp.usage_breakdown_list = vec![
221            UsageBreakdown {
222                resource_type: None,
223                current_usage_with_precision: Some(1.0),
224                usage_limit_with_precision: Some(2.0),
225            },
226            UsageBreakdown {
227                resource_type: Some("OTHER".into()),
228                current_usage_with_precision: Some(3.0),
229                usage_limit_with_precision: Some(4.0),
230            },
231        ];
232        assert!(matches!(to_snapshot(resp), Err(AppError::Schema(_))));
233    }
234
235    #[test]
236    fn missing_reset_is_none_not_an_error() {
237        let mut resp = sample();
238        resp.next_date_reset = None;
239        let snap = to_snapshot(resp).unwrap();
240        assert_eq!(snap.reset_at, None);
241    }
242
243    #[test]
244    fn missing_plan_is_schema_drift() {
245        let mut resp = sample();
246        resp.subscription_info = None;
247        assert!(matches!(to_snapshot(resp), Err(AppError::Schema(_))));
248    }
249
250    #[test]
251    fn empty_breakdown_list_is_schema_drift() {
252        let mut resp = sample();
253        resp.usage_breakdown_list.clear();
254        assert!(matches!(to_snapshot(resp), Err(AppError::Schema(_))));
255    }
256
257    #[test]
258    fn non_finite_usage_is_schema_drift() {
259        let mut resp = sample();
260        resp.usage_breakdown_list[0].current_usage_with_precision = Some(f64::NAN);
261        assert!(matches!(to_snapshot(resp), Err(AppError::Schema(_))));
262    }
263
264    #[test]
265    fn negative_usage_is_schema_drift() {
266        let mut resp = sample();
267        resp.usage_breakdown_list[0].current_usage_with_precision = Some(-1.0);
268        assert!(matches!(to_snapshot(resp), Err(AppError::Schema(_))));
269    }
270}