Skip to main content

usage_monitor_cli/provider/
anthropic.rs

1use async_trait::async_trait;
2
3use crate::error::SpendPanelError;
4use crate::model::{
5    CostSnapshot, DailyCost, PlanInfo, RateWindow, RateWindowStatus, SpendLimit, UsageSnapshot,
6};
7use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
8
9// ---------------------------------------------------------------------------
10// Helper types for deserializing Anthropic responses
11// ---------------------------------------------------------------------------
12
13#[derive(serde::Deserialize, Debug)]
14struct UsageReportItem {
15    date: Option<String>,
16    #[allow(dead_code)]
17    model: Option<String>,
18    input_tokens: Option<u64>,
19    output_tokens: Option<u64>,
20}
21
22#[derive(serde::Deserialize, Debug)]
23struct UsageReportResponse {
24    data: Vec<UsageReportItem>,
25}
26
27#[derive(serde::Deserialize, Debug)]
28struct CostReportItem {
29    date: Option<String>,
30    cost: Option<CostValue>,
31}
32
33#[derive(serde::Deserialize, Debug)]
34struct CostValue {
35    value: String,
36    #[allow(dead_code)]
37    currency: String,
38}
39
40#[derive(serde::Deserialize, Debug)]
41struct CostReportResponse {
42    data: Vec<CostReportItem>,
43}
44
45// ---------------------------------------------------------------------------
46// Provider
47// ---------------------------------------------------------------------------
48
49pub struct AnthropicProvider {
50    metadata: ProviderMetadata,
51    /// Base URL override for tests.
52    base_url: Option<String>,
53}
54
55impl AnthropicProvider {
56    pub fn new() -> Self {
57        Self {
58            metadata: ProviderMetadata {
59                id: "anthropic",
60                name: "Anthropic (Claude)",
61                description: "Anthropic Claude API usage monitor",
62                auth_methods: &["api_key", "oauth", "cli"],
63                website: Some("https://docs.anthropic.com"),
64            },
65            base_url: None,
66        }
67    }
68
69    /// Creates a provider with a custom base URL (for tests).
70    pub fn with_base_url(url: &str) -> Self {
71        let mut p = Self::new();
72        p.base_url = Some(url.to_string());
73        p
74    }
75
76    fn api_base(&self) -> &str {
77        self.base_url
78            .as_deref()
79            .unwrap_or("https://api.anthropic.com")
80    }
81
82    /// Detection helper: a non-empty API key is available.
83    fn detect_credentials_from(key: Option<&str>) -> bool {
84        key.is_some_and(|k| !k.is_empty())
85    }
86
87    fn resolve_api_key(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
88        if let Some(key) = ctx.config.get("api_key")
89            && !key.is_empty()
90        {
91            return Ok(key.clone());
92        }
93        if let Ok(key) = std::env::var("ANTHROPIC_API_KEY")
94            && !key.is_empty()
95        {
96            return Ok(key);
97        }
98        Err(SpendPanelError::AuthFailed(
99            "anthropic".into(),
100            "no API key in config or ANTHROPIC_API_KEY env var".into(),
101        ))
102    }
103
104    fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
105        reqwest::Client::builder()
106            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
107            .build()
108            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
109    }
110
111    /// Fetch /v1/organizations/usage_report/messages
112    async fn fetch_usage_report(
113        base_url: &str,
114        client: &reqwest::Client,
115        api_key: &str,
116    ) -> Result<Vec<UsageReportItem>, SpendPanelError> {
117        let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
118        let week_ago = (chrono::Utc::now() - chrono::Duration::days(7))
119            .format("%Y-%m-%d")
120            .to_string();
121
122        let url = format!("{}/v1/organizations/usage_report/messages", base_url);
123        let resp = client
124            .get(&url)
125            .header("x-api-key", api_key)
126            .header("anthropic-version", "2023-06-01")
127            .query(&[
128                ("start_date", &*week_ago),
129                ("end_date", &*today),
130                ("bucket_width", "1d"),
131            ])
132            .send()
133            .await
134            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
135
136        let status = resp.status();
137        let body = resp
138            .text()
139            .await
140            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
141
142        if status == 401 {
143            return Err(SpendPanelError::AuthFailed(
144                "anthropic".into(),
145                "invalid API key".into(),
146            ));
147        }
148        if !status.is_success() {
149            return Err(SpendPanelError::ProviderError(
150                "anthropic".into(),
151                format!("HTTP {}: {}", status, body),
152            ));
153        }
154
155        let report: UsageReportResponse = serde_json::from_str(&body).map_err(|e| {
156            SpendPanelError::ParseError("anthropic".into(), format!("usage report: {}", e))
157        })?;
158        Ok(report.data)
159    }
160
161    /// Fetch /v1/organizations/cost_report
162    async fn fetch_cost_report(
163        base_url: &str,
164        client: &reqwest::Client,
165        api_key: &str,
166    ) -> Result<Vec<CostReportItem>, SpendPanelError> {
167        let today = chrono::Utc::now().format("%Y-%m-%d").to_string();
168        let week_ago = (chrono::Utc::now() - chrono::Duration::days(7))
169            .format("%Y-%m-%d")
170            .to_string();
171
172        let url = format!("{}/v1/organizations/cost_report", base_url);
173        let resp = client
174            .get(&url)
175            .header("x-api-key", api_key)
176            .header("anthropic-version", "2023-06-01")
177            .query(&[("start_date", &*week_ago), ("end_date", &*today)])
178            .send()
179            .await
180            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
181
182        if !resp.status().is_success() {
183            return Ok(Vec::new());
184        }
185
186        let body = resp
187            .text()
188            .await
189            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
190        let report: CostReportResponse = serde_json::from_str(&body).map_err(|e| {
191            SpendPanelError::ParseError("anthropic".into(), format!("cost report: {}", e))
192        })?;
193        Ok(report.data)
194    }
195
196    /// Probes POST /v1/messages to capture rate limit headers.
197    async fn probe_rate_limits(
198        base_url: &str,
199        client: &reqwest::Client,
200        api_key: &str,
201    ) -> Result<Option<RateWindow>, SpendPanelError> {
202        let url = format!("{}/v1/messages", base_url);
203        let resp = client
204            .post(&url)
205            .header("x-api-key", api_key)
206            .header("anthropic-version", "2023-06-01")
207            .header("content-type", "application/json")
208            .body(r#"{"model":"claude-sonnet-4-20250514","max_tokens":1,"messages":[{"role":"user","content":"hi"}]}"#)
209            .send()
210            .await
211            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
212
213        let headers = resp.headers();
214
215        if let (Some(limit_str), Some(rem_str)) = (
216            headers
217                .get("x-ratelimit-limit-requests")
218                .and_then(|v| v.to_str().ok()),
219            headers
220                .get("x-ratelimit-remaining-requests")
221                .and_then(|v| v.to_str().ok()),
222        ) && let (Ok(limit), Ok(remaining)) = (limit_str.parse::<u64>(), rem_str.parse::<u64>())
223        {
224            let used = limit.saturating_sub(remaining);
225            return Ok(Some(RateWindow::new(used, limit, "RPM", 1)));
226        }
227
228        Ok(None)
229    }
230}
231
232impl Default for AnthropicProvider {
233    fn default() -> Self {
234        Self::new()
235    }
236}
237
238#[async_trait]
239impl UsageProvider for AnthropicProvider {
240    fn metadata(&self) -> &ProviderMetadata {
241        &self.metadata
242    }
243
244    fn detect_credentials(&self) -> bool {
245        AnthropicProvider::detect_credentials_from(
246            std::env::var("ANTHROPIC_API_KEY").ok().as_deref(),
247        )
248    }
249
250    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
251        let api_key = Self::resolve_api_key(ctx)?;
252        let client = Self::build_client(ctx)?;
253        let base = self.api_base();
254
255        let (rate_limit_res, usage_res, cost_res) = tokio::join!(
256            Self::probe_rate_limits(base, &client, &api_key),
257            Self::fetch_usage_report(base, &client, &api_key),
258            Self::fetch_cost_report(base, &client, &api_key),
259        );
260
261        let mut snapshot = UsageSnapshot::new("anthropic");
262        snapshot.collected_at = chrono::Utc::now();
263        let mut usage_ok = false;
264
265        // Rate limit
266        if let Ok(Some(rl)) = rate_limit_res {
267            snapshot.primary_rate_window = Some(rl);
268        }
269
270        // Usage report
271        if let Ok(ref items) = usage_res {
272            usage_ok = true;
273            let mut total_in = 0u64;
274            let mut total_out = 0u64;
275            let mut daily: std::collections::HashMap<String, (u64, u64)> =
276                std::collections::HashMap::new();
277
278            for item in items {
279                total_in += item.input_tokens.unwrap_or(0);
280                total_out += item.output_tokens.unwrap_or(0);
281                let date = item.date.clone().unwrap_or_default();
282                let e = daily.entry(date).or_default();
283                e.0 += item.input_tokens.unwrap_or(0);
284                e.1 += item.output_tokens.unwrap_or(0);
285            }
286
287            if total_in + total_out > 0 {
288                snapshot.secondary_rate_window = Some(RateWindow {
289                    label: "7-day Token Usage".into(),
290                    window_minutes: 10080,
291                    usage_ratio: 0.0,
292                    limit: None,
293                    used: Some(total_in + total_out),
294                    remaining: None,
295                    resets_at: None,
296                    status: RateWindowStatus::Normal,
297                });
298            }
299
300            let daily_costs: Vec<DailyCost> = daily
301                .into_iter()
302                .filter_map(|(ds, (inp, out))| {
303                    let date = chrono::NaiveDate::parse_from_str(&ds, "%Y-%m-%d").ok()?;
304                    Some(DailyCost {
305                        date,
306                        cost: 0.0,
307                        tokens_input: Some(inp),
308                        tokens_output: Some(out),
309                        requests: None,
310                    })
311                })
312                .collect();
313
314            if !daily_costs.is_empty() {
315                snapshot.cost = Some(CostSnapshot {
316                    total_cost: None,
317                    currency: "USD".into(),
318                    daily_costs,
319                    spend_limit: None,
320                });
321            }
322        }
323
324        // Cost report
325        if let Ok(cost_items) = cost_res
326            && !cost_items.is_empty()
327        {
328            let mut daily: Vec<DailyCost> = cost_items
329                .into_iter()
330                .filter_map(|item| {
331                    let date = chrono::NaiveDate::parse_from_str(
332                        &item.date.unwrap_or_default(),
333                        "%Y-%m-%d",
334                    )
335                    .ok()?;
336                    let cost = item
337                        .cost
338                        .as_ref()
339                        .and_then(|c| c.value.parse::<f64>().ok())
340                        .unwrap_or(0.0);
341                    Some(DailyCost {
342                        date,
343                        cost,
344                        tokens_input: None,
345                        tokens_output: None,
346                        requests: None,
347                    })
348                })
349                .collect();
350            daily.sort_by_key(|a| a.date);
351
352            let total: f64 = daily.iter().map(|d| d.cost).sum();
353            snapshot.cost = Some(CostSnapshot {
354                total_cost: Some(total),
355                currency: "USD".into(),
356                daily_costs: daily,
357                spend_limit: Some(SpendLimit {
358                    limit: 50.0,
359                    used: total,
360                    period: "monthly".into(),
361                }),
362            });
363            snapshot.plan = Some(PlanInfo {
364                name: "API".into(),
365                tier: None,
366                features: vec![],
367                price: None,
368                currency: Some("USD".into()),
369                billing_period: Some("monthly".into()),
370            });
371        }
372
373        if snapshot.plan.is_none() && usage_ok {
374            snapshot.plan = Some(PlanInfo {
375                name: "API".into(),
376                tier: None,
377                features: vec![],
378                price: None,
379                currency: Some("USD".into()),
380                billing_period: Some("monthly".into()),
381            });
382        }
383
384        Ok(snapshot)
385    }
386}
387
388// ---------------------------------------------------------------------------
389// Tests
390// ---------------------------------------------------------------------------
391
392#[cfg(test)]
393mod tests {
394    use super::*;
395    use wiremock::matchers::{header, method, path};
396    use wiremock::{Mock, MockServer, ResponseTemplate};
397
398    #[test]
399    fn test_resolve_api_key_from_context() {
400        let mut ctx = ProviderContext::new();
401        ctx.config.insert("api_key".into(), "sk-ant-test".into());
402        assert_eq!(
403            AnthropicProvider::resolve_api_key(&ctx).unwrap(),
404            "sk-ant-test"
405        );
406    }
407
408    #[test]
409    fn test_resolve_api_key_missing_is_error() {
410        let ctx = ProviderContext::new();
411        assert!(AnthropicProvider::resolve_api_key(&ctx).is_err());
412    }
413
414    #[test]
415    fn test_provider_metadata() {
416        let p = AnthropicProvider::new();
417        let m = p.metadata();
418        assert_eq!(m.id, "anthropic");
419        assert!(m.auth_methods.contains(&"api_key"));
420    }
421
422    #[test]
423    fn test_detect_credentials_from() {
424        assert!(AnthropicProvider::detect_credentials_from(Some("sk-ant-x")));
425        assert!(!AnthropicProvider::detect_credentials_from(Some("")));
426        assert!(!AnthropicProvider::detect_credentials_from(None));
427    }
428
429    #[tokio::test]
430    async fn test_fetch_usage_report_success() {
431        let server = MockServer::start().await;
432
433        Mock::given(method("GET"))
434            .and(path("/v1/organizations/usage_report/messages"))
435            .and(header("x-api-key", "sk-ant-test"))
436            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
437                "data": [
438                    {"date": "2026-06-12", "model": "claude-sonnet-4", "input_tokens": 250000, "output_tokens": 75000},
439                    {"date": "2026-06-11", "model": "claude-sonnet-4", "input_tokens": 120000, "output_tokens": 35000},
440                ]
441            })))
442            .mount(&server)
443            .await;
444
445        let client = reqwest::Client::new();
446        let items =
447            AnthropicProvider::fetch_usage_report(&server.uri(), &client, "sk-ant-test").await;
448        assert!(items.is_ok());
449        assert_eq!(items.unwrap().len(), 2);
450    }
451
452    #[tokio::test]
453    async fn test_fetch_usage_report_401() {
454        let server = MockServer::start().await;
455        Mock::given(method("GET"))
456            .and(path("/v1/organizations/usage_report/messages"))
457            .respond_with(ResponseTemplate::new(401))
458            .mount(&server)
459            .await;
460
461        let client = reqwest::Client::new();
462        let result = AnthropicProvider::fetch_usage_report(&server.uri(), &client, "bad").await;
463        assert!(matches!(result, Err(SpendPanelError::AuthFailed(_, _))));
464    }
465
466    #[tokio::test]
467    async fn test_probe_rate_limits_parses_headers() {
468        let server = MockServer::start().await;
469
470        Mock::given(method("POST"))
471            .and(path("/v1/messages"))
472            .respond_with(
473                ResponseTemplate::new(200)
474                    .set_body_json(serde_json::json!({"id":"m","type":"message","content":[]}))
475                    .insert_header("x-ratelimit-limit-requests", "100")
476                    .insert_header("x-ratelimit-remaining-requests", "72"),
477            )
478            .mount(&server)
479            .await;
480
481        let client = reqwest::Client::new();
482        let result =
483            AnthropicProvider::probe_rate_limits(&server.uri(), &client, "sk-ant-test").await;
484        assert!(result.is_ok());
485        let rl = result.unwrap().unwrap();
486        assert_eq!(rl.used, Some(28));
487        assert!((rl.usage_ratio - 0.28).abs() < f64::EPSILON);
488    }
489
490    #[tokio::test]
491    async fn test_probe_rate_limits_no_headers() {
492        let server = MockServer::start().await;
493        Mock::given(method("POST"))
494            .and(path("/v1/messages"))
495            .respond_with(
496                ResponseTemplate::new(200)
497                    .set_body_json(serde_json::json!({"id":"m","type":"message","content":[]})),
498            )
499            .mount(&server)
500            .await;
501
502        let client = reqwest::Client::new();
503        let result = AnthropicProvider::probe_rate_limits(&server.uri(), &client, "key").await;
504        assert!(result.is_ok());
505        assert!(result.unwrap().is_none());
506    }
507
508    #[tokio::test]
509    async fn test_fetch_cost_report_success() {
510        let server = MockServer::start().await;
511
512        Mock::given(method("GET"))
513            .and(path("/v1/organizations/cost_report"))
514            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
515                "data": [{"date": "2026-06-12", "cost": {"value": "3.50", "currency": "USD"}}]
516            })))
517            .mount(&server)
518            .await;
519
520        let client = reqwest::Client::new();
521        let items =
522            AnthropicProvider::fetch_cost_report(&server.uri(), &client, "sk-ant-admin").await;
523        assert!(items.is_ok());
524        assert_eq!(items.unwrap().len(), 1);
525    }
526
527    #[tokio::test]
528    async fn test_fetch_cost_report_non_admin_returns_empty() {
529        let server = MockServer::start().await;
530        Mock::given(method("GET"))
531            .and(path("/v1/organizations/cost_report"))
532            .respond_with(ResponseTemplate::new(403))
533            .mount(&server)
534            .await;
535
536        let client = reqwest::Client::new();
537        let items = AnthropicProvider::fetch_cost_report(&server.uri(), &client, "non-admin").await;
538        assert!(items.is_ok());
539        assert!(items.unwrap().is_empty());
540    }
541
542    #[tokio::test]
543    async fn test_integrated_fetch_with_mocks() {
544        let server = MockServer::start().await;
545
546        // Mock usage report
547        Mock::given(method("GET"))
548            .and(path("/v1/organizations/usage_report/messages"))
549            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
550                "data": [{"date": "2026-06-12", "model": "claude-sonnet-4", "input_tokens": 100000, "output_tokens": 30000}]
551            })))
552            .mount(&server)
553            .await;
554
555        // Mock probe
556        Mock::given(method("POST"))
557            .and(path("/v1/messages"))
558            .respond_with(
559                ResponseTemplate::new(200)
560                    .set_body_json(serde_json::json!({"id":"m","type":"message","content":[]}))
561                    .insert_header("x-ratelimit-limit-requests", "50")
562                    .insert_header("x-ratelimit-remaining-requests", "40"),
563            )
564            .mount(&server)
565            .await;
566
567        // Mock cost report
568        Mock::given(method("GET"))
569            .and(path("/v1/organizations/cost_report"))
570            .respond_with(ResponseTemplate::new(200).set_body_json(serde_json::json!({
571                "data": [{"date": "2026-06-12", "cost": {"value": "2.50", "currency": "USD"}}]
572            })))
573            .mount(&server)
574            .await;
575
576        let provider = AnthropicProvider::with_base_url(&server.uri());
577        let mut ctx = ProviderContext::new();
578        ctx.config.insert("api_key".into(), "sk-ant-test".into());
579
580        let result = provider.fetch_usage(&ctx).await;
581        assert!(result.is_ok());
582        let snap = result.unwrap();
583
584        assert_eq!(snap.provider_id, "anthropic");
585        assert!(snap.primary_rate_window.is_some());
586        assert_eq!(snap.primary_rate_window.unwrap().used, Some(10)); // 50-40
587
588        assert!(snap.cost.is_some());
589        let cost = snap.cost.unwrap();
590        assert!(cost.total_cost.is_some());
591        assert!((cost.total_cost.unwrap() - 2.50).abs() < f64::EPSILON);
592    }
593}