use anyhow::{Context, Result};
use serde::{Deserialize, Serialize};
use serde_json::Value;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RealTimeQuota {
pub account_id: String,
pub plan: String,
pub usage_this_month: u64,
pub quota_limit: u64,
pub remaining_quota: u64,
pub percent_used: f64,
pub reset_date: Option<String>,
}
impl RealTimeQuota {
pub fn days_until_reset(&self) -> Option<i64> {
use chrono::{DateTime, Utc};
self.reset_date.as_ref().and_then(|date| {
DateTime::parse_from_rfc3339(date)
.ok()
.map(|d| (d.with_timezone(&Utc) - Utc::now()).num_days())
})
}
pub fn is_critical(&self) -> bool {
self.percent_used > 80.0
}
pub fn is_low(&self) -> bool {
self.percent_used > 50.0
}
}
pub async fn fetch_quota(api_key: &str) -> Result<RealTimeQuota> {
let client = reqwest::Client::new();
let response = client
.get("https://api.openai.com/v1/dashboard/billing/subscription")
.header("Authorization", format!("Bearer {api_key}"))
.send()
.await
.context("Failed to connect to OpenAI API")?;
if !response.status().is_success() {
let status = response.status();
let text = response.text().await.unwrap_or_default();
anyhow::bail!("OpenAI API error ({status}): {text}");
}
let data: serde_json::Value = response
.json()
.await
.context("Failed to parse OpenAI API response")?;
let account_id = data
.get("account_id")
.and_then(Value::as_str)
.unwrap_or("unknown")
.to_string();
let plan = data
.get("plan")
.and_then(Value::as_str)
.unwrap_or("unknown")
.to_string();
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
let quota_limit = data
.get("hard_limit_usd")
.and_then(Value::as_f64)
.map(|v| (v * 100.0) as u64)
.map_or(0, std::convert::identity);
let usage = fetch_usage(api_key).await.map_or(0, std::convert::identity);
let remaining = quota_limit.saturating_sub(usage);
#[allow(clippy::cast_precision_loss)]
let percent_used = if quota_limit > 0 {
(usage as f64 / quota_limit as f64) * 100.0
} else {
0.0
};
let reset_date = data
.get("reset_date")
.or_else(|| data.get("billing_cycle_anchor"))
.and_then(Value::as_str)
.map(std::string::ToString::to_string);
Ok(RealTimeQuota {
account_id,
plan,
usage_this_month: usage,
quota_limit,
remaining_quota: remaining,
percent_used,
reset_date,
})
}
async fn fetch_usage(api_key: &str) -> Result<u64> {
use chrono::{Datelike, Utc};
let now = Utc::now();
let start_of_month = format!("{}-{:02}-01", now.year(), now.month());
let today = format!("{}-{:02}-{:02}", now.year(), now.month(), now.day());
let client = reqwest::Client::new();
let url = format!(
"https://api.openai.com/v1/dashboard/billing/usage?start_date={start_of_month}&end_date={today}"
);
let response = client
.get(&url)
.header("Authorization", format!("Bearer {api_key}"))
.send()
.await
.context("Failed to fetch usage data")?;
if !response.status().is_success() {
let status = response.status();
let text = response.text().await.unwrap_or_default();
anyhow::bail!("OpenAI API error ({status}): {text}");
}
let data: serde_json::Value = response
.json()
.await
.context("Failed to parse usage response")?;
#[allow(clippy::cast_possible_truncation, clippy::cast_sign_loss)]
let total_usage = data
.get("total_usage")
.and_then(Value::as_f64)
.map(|v| (v * 100.0) as u64) .map_or(0, std::convert::identity);
Ok(total_usage)
}
pub fn extract_api_key(auth_json: &serde_json::Value) -> Option<String> {
auth_json
.get("api_key")
.or_else(|| auth_json.get("key"))
.or_else(|| auth_json.get("access_token"))
.and_then(Value::as_str)
.map(std::string::ToString::to_string)
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn test_realtime_quota_calculations() {
let quota = RealTimeQuota {
account_id: "test".to_string(),
plan: "personal".to_string(),
usage_this_month: 7500, quota_limit: 10000, remaining_quota: 2500, percent_used: 75.0,
reset_date: None,
};
assert!(quota.is_low());
assert!(!quota.is_critical());
assert_eq!(quota.remaining_quota, 2500);
}
#[test]
fn test_extract_api_key() {
let auth = serde_json::json!({
"api_key": "sk-test123",
"email": "test@example.com"
});
assert_eq!(extract_api_key(&auth), Some("sk-test123".to_string()));
}
#[test]
fn test_extract_api_key_alt_field() {
let auth = serde_json::json!({
"access_token": "sk-token456",
});
assert_eq!(extract_api_key(&auth), Some("sk-token456".to_string()));
}
}