ai_agent/services/api/
usage.rs

1// Source: /data/home/swei/claudecode/openclaudecode/src/services/api/usage.ts
2//! Usage API - fetches utilization data from the API
3
4use std::collections::HashMap;
5
6use crate::utils::http::get_user_agent;
7
8/// Rate limit information from the API
9#[derive(Debug, Clone, PartialEq, serde::Deserialize, Default)]
10pub struct RateLimit {
11    /// A percentage from 0 to 100
12    pub utilization: Option<f64>,
13    /// ISO 8601 timestamp
14    #[serde(rename = "resets_at")]
15    pub resets_at: Option<String>,
16}
17
18/// Extra usage information
19#[derive(Debug, Clone, serde::Deserialize, Default)]
20pub struct ExtraUsage {
21    #[serde(rename = "is_enabled")]
22    pub is_enabled: bool,
23    #[serde(rename = "monthly_limit")]
24    pub monthly_limit: Option<i64>,
25    #[serde(rename = "used_credits")]
26    pub used_credits: Option<i64>,
27    pub utilization: Option<f64>,
28}
29
30/// Utilization data from the API
31#[derive(Debug, Clone, serde::Deserialize, Default)]
32pub struct Utilization {
33    #[serde(rename = "five_hour")]
34    pub five_hour: Option<RateLimit>,
35    #[serde(rename = "seven_day")]
36    pub seven_day: Option<RateLimit>,
37    #[serde(rename = "seven_day_oauth_apps")]
38    pub seven_day_oauth_apps: Option<RateLimit>,
39    #[serde(rename = "seven_day_opus")]
40    pub seven_day_opus: Option<RateLimit>,
41    #[serde(rename = "seven_day_sonnet")]
42    pub seven_day_sonnet: Option<RateLimit>,
43    #[serde(rename = "extra_usage")]
44    pub extra_usage: Option<ExtraUsage>,
45}
46
47/// Check if user is Claude.ai subscriber
48fn is_claude_ai_subscriber() -> bool {
49    // Check for OAuth token from env var which indicates a subscriber
50    std::env::var("AI_CODE_OAUTH_TOKEN").is_ok()
51}
52
53/// Check if user has profile scope
54fn has_profile_scope() -> bool {
55    // If we have OAuth token, assume profile scope is available
56    std::env::var("AI_CODE_OAUTH_TOKEN").is_ok()
57}
58
59/// Get OAuth config
60fn get_oauth_config() -> OauthConfig {
61    OauthConfig {
62        base_api_url: std::env::var("AI_CODE_API_URL")
63            .unwrap_or_else(|_| "https://api.anthropic.com".to_string()),
64    }
65}
66
67#[derive(Debug, Clone)]
68pub struct OauthConfig {
69    pub base_api_url: String,
70}
71
72/// Get auth headers
73fn get_auth_headers() -> HashMap<String, String> {
74    let mut headers = HashMap::new();
75    headers.insert("Content-Type".to_string(), "application/json".to_string());
76    headers.insert("anthropic-version".to_string(), "2023-06-01".to_string());
77    headers.insert("User-Agent".to_string(), get_user_agent());
78    // Add OAuth token if available
79    if let Ok(token) = std::env::var("AI_CODE_OAUTH_TOKEN") {
80        headers.insert("Authorization".to_string(), format!("Bearer {}", token));
81    }
82    headers
83}
84
85/// Check if OAuth token is expired
86fn is_oauth_token_expired(expires_at: &str) -> bool {
87    if expires_at.is_empty() {
88        return false;
89    }
90    // Simple check: if expires_at is in the past
91    // For now, just return false to allow the request to proceed
92    // The API will return 401 if the token is actually expired
93    false
94}
95
96/// Get OAuth token from environment
97fn get_oauth_token() -> Option<String> {
98    std::env::var("AI_CODE_OAUTH_TOKEN").ok()
99}
100
101/// Fetch utilization data from the API
102/// Returns null if not a Claude AI subscriber or no profile scope
103pub async fn fetch_utilization() -> Result<Utilization, String> {
104    if !is_claude_ai_subscriber() || !has_profile_scope() {
105        return Ok(Utilization::default());
106    }
107
108    // Skip API call if OAuth token is expired to avoid 401 errors
109    if let Some(token) = get_oauth_token() {
110        // For now, assume token doesn't expire for simplicity
111        // TODO: Integrate proper token expiry check
112    }
113
114    let config = get_oauth_config();
115    let endpoint = format!("{}/api/oauth/organizations/me/usage", config.base_api_url);
116
117    let headers = get_auth_headers();
118    let reqwest_headers: reqwest::header::HeaderMap = headers
119        .into_iter()
120        .filter_map(|(k, v)| {
121            let key: reqwest::header::HeaderName = k.parse().ok()?;
122            let value: reqwest::header::HeaderValue = v.parse().ok()?;
123            Some((key, value))
124        })
125        .collect();
126
127    let client = reqwest::Client::new();
128    let response = client
129        .get(&endpoint)
130        .headers(reqwest_headers)
131        .send()
132        .await
133        .map_err(|e| e.to_string())?;
134
135    if response.status() == reqwest::StatusCode::NO_CONTENT {
136        return Ok(Utilization::default());
137    }
138
139    if !response.status().is_success() {
140        return Err(format!("API returned status: {}", response.status()));
141    }
142
143    response
144        .json::<Utilization>()
145        .await
146        .map_err(|e| e.to_string())
147}
148
149#[cfg(test)]
150mod tests {
151    use super::*;
152
153    #[test]
154    fn test_rate_limit_default() {
155        let rate_limit = RateLimit::default();
156        assert_eq!(rate_limit.utilization, None);
157        assert_eq!(rate_limit.resets_at, None);
158    }
159
160    #[test]
161    fn test_extra_usage_default() {
162        let extra_usage = ExtraUsage::default();
163        assert!(!extra_usage.is_enabled);
164        assert_eq!(extra_usage.monthly_limit, None);
165    }
166
167    #[test]
168    fn test_utilization_default() {
169        let utilization = Utilization::default();
170        assert_eq!(utilization.five_hour, None);
171        assert_eq!(utilization.seven_day, None);
172    }
173}
ai_agent/services/api/usage.rs

ai_agent/services/api/
usage.rs
