Skip to main content

usage_monitor_cli/provider/
openrouter.rs

1use async_trait::async_trait;
2
3use crate::error::SpendPanelError;
4use crate::model::{
5    CostSnapshot, CreditsSnapshot, RateWindow, RateWindowStatus, SpendLimit, UsageSnapshot,
6};
7use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
8
9#[derive(Debug, serde::Deserialize)]
10struct OpenRouterCreditsResponse {
11    data: OpenRouterCreditsData,
12}
13
14#[derive(Debug, serde::Deserialize)]
15struct OpenRouterCreditsData {
16    total_credits: f64,
17    total_usage: f64,
18}
19
20impl OpenRouterCreditsData {
21    fn balance(&self) -> f64 {
22        (self.total_credits - self.total_usage).max(0.0)
23    }
24}
25
26#[derive(Debug, serde::Deserialize)]
27struct OpenRouterKeyResponse {
28    data: OpenRouterKeyData,
29}
30
31#[derive(Debug, serde::Deserialize, Clone, PartialEq)]
32struct OpenRouterKeyData {
33    limit: Option<f64>,
34    usage: Option<f64>,
35    usage_daily: Option<f64>,
36    usage_weekly: Option<f64>,
37    usage_monthly: Option<f64>,
38    rate_limit: Option<OpenRouterRateLimit>,
39}
40
41#[derive(Debug, serde::Deserialize, Clone, PartialEq)]
42struct OpenRouterRateLimit {
43    requests: u64,
44    interval: String,
45}
46
47#[derive(Debug, Clone, PartialEq)]
48struct OpenRouterUsage {
49    total_credits: f64,
50    total_usage: f64,
51    balance: f64,
52    key_data_fetched: bool,
53    key_limit: Option<f64>,
54    key_usage: Option<f64>,
55    key_usage_daily: Option<f64>,
56    key_usage_weekly: Option<f64>,
57    key_usage_monthly: Option<f64>,
58    rate_limit: Option<OpenRouterRateLimit>,
59}
60
61impl OpenRouterUsage {
62    fn has_valid_key_quota(&self) -> bool {
63        matches!((self.key_limit, self.key_usage), (Some(limit), Some(usage)) if limit > 0.0 && usage >= 0.0)
64    }
65
66    fn key_used_ratio(&self) -> Option<f64> {
67        if !self.has_valid_key_quota() {
68            return None;
69        }
70        Some((self.key_usage.unwrap() / self.key_limit.unwrap()).clamp(0.0, 1.0))
71    }
72}
73
74/// OpenRouter API credits provider.
75pub struct OpenRouterProvider {
76    metadata: ProviderMetadata,
77    /// Base URL override for tests.
78    base_url: Option<String>,
79}
80
81impl OpenRouterProvider {
82    pub fn new() -> Self {
83        Self {
84            metadata: ProviderMetadata {
85                id: "openrouter",
86                name: "OpenRouter",
87                description: "OpenRouter credits and API-key usage monitor",
88                auth_methods: &["api_key", "env"],
89                website: Some("https://openrouter.ai"),
90            },
91            base_url: None,
92        }
93    }
94
95    /// Creates a provider with a custom base URL (for tests).
96    pub fn with_base_url(url: &str) -> Self {
97        let mut p = Self::new();
98        p.base_url = Some(url.to_string());
99        p
100    }
101
102    fn clean(raw: &str) -> String {
103        let mut value = raw.trim();
104        if value.len() >= 2
105            && ((value.starts_with('"') && value.ends_with('"'))
106                || (value.starts_with('\'') && value.ends_with('\'')))
107        {
108            value = &value[1..value.len() - 1];
109        }
110        value.trim().to_string()
111    }
112
113    fn detect_credentials_from(key: Option<&str>) -> bool {
114        key.map(Self::clean).is_some_and(|key| !key.is_empty())
115    }
116
117    fn resolve_api_key(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
118        for key in ["api_key", "token"] {
119            if let Some(value) = ctx.config.get(key) {
120                let cleaned = Self::clean(value);
121                if !cleaned.is_empty() {
122                    return Ok(cleaned);
123                }
124            }
125        }
126        if let Ok(value) = std::env::var("OPENROUTER_API_KEY") {
127            let cleaned = Self::clean(&value);
128            if !cleaned.is_empty() {
129                return Ok(cleaned);
130            }
131        }
132        Err(SpendPanelError::AuthFailed(
133            "openrouter".into(),
134            "no API key found in config, token, or OPENROUTER_API_KEY".into(),
135        ))
136    }
137
138    fn api_base(&self, ctx: &ProviderContext) -> String {
139        let configured = ctx
140            .config
141            .get("api_url")
142            .or_else(|| ctx.config.get("base_url"))
143            .map(String::as_str)
144            .filter(|value| !value.is_empty())
145            .map(Self::clean)
146            .or_else(|| {
147                std::env::var("OPENROUTER_API_URL")
148                    .ok()
149                    .map(|v| Self::clean(&v))
150            })
151            .or_else(|| self.base_url.clone())
152            .unwrap_or_else(|| "https://openrouter.ai/api/v1".into());
153
154        if configured.starts_with("http://") || configured.starts_with("https://") {
155            configured.trim_end_matches('/').to_string()
156        } else {
157            format!("https://{}", configured.trim_end_matches('/'))
158        }
159    }
160
161    fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
162        reqwest::Client::builder()
163            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
164            .build()
165            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
166    }
167
168    fn add_headers(
169        builder: reqwest::RequestBuilder,
170        api_key: &str,
171        ctx: &ProviderContext,
172    ) -> reqwest::RequestBuilder {
173        let builder = builder
174            .header("Authorization", format!("Bearer {}", api_key))
175            .header("Accept", "application/json");
176
177        let referer = ctx
178            .config
179            .get("http_referer")
180            .map(|v| Self::clean(v))
181            .filter(|v| !v.is_empty())
182            .or_else(|| {
183                std::env::var("OPENROUTER_HTTP_REFERER")
184                    .ok()
185                    .map(|v| Self::clean(&v))
186            });
187        let builder = if let Some(referer) = referer {
188            builder.header("HTTP-Referer", referer)
189        } else {
190            builder
191        };
192
193        let title = ctx
194            .config
195            .get("x_title")
196            .map(|v| Self::clean(v))
197            .filter(|v| !v.is_empty())
198            .or_else(|| {
199                std::env::var("OPENROUTER_X_TITLE")
200                    .ok()
201                    .map(|v| Self::clean(&v))
202            })
203            .unwrap_or_else(|| "UsageMonitor".into());
204        builder.header("X-Title", title)
205    }
206
207    async fn get_json<T: serde::de::DeserializeOwned>(
208        client: &reqwest::Client,
209        url: String,
210        api_key: &str,
211        ctx: &ProviderContext,
212    ) -> Result<T, SpendPanelError> {
213        let resp = Self::add_headers(client.get(url), api_key, ctx)
214            .send()
215            .await
216            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
217        let status = resp.status();
218        let body = resp
219            .text()
220            .await
221            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
222
223        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
224            return Err(SpendPanelError::AuthFailed(
225                "openrouter".into(),
226                format!("invalid API key (HTTP {})", status.as_u16()),
227            ));
228        }
229        if !status.is_success() {
230            return Err(SpendPanelError::ProviderError(
231                "openrouter".into(),
232                format!("HTTP {}", status),
233            ));
234        }
235
236        serde_json::from_str(&body)
237            .map_err(|e| SpendPanelError::ParseError("openrouter".into(), e.to_string()))
238    }
239
240    async fn fetch_key_data(
241        client: &reqwest::Client,
242        base_url: &str,
243        api_key: &str,
244        ctx: &ProviderContext,
245    ) -> Option<OpenRouterKeyData> {
246        let result = Self::get_json::<OpenRouterKeyResponse>(
247            client,
248            format!("{}/key", base_url),
249            api_key,
250            ctx,
251        )
252        .await;
253        result.ok().map(|response| response.data)
254    }
255
256    async fn fetch_usage_data(
257        &self,
258        client: &reqwest::Client,
259        api_key: &str,
260        ctx: &ProviderContext,
261    ) -> Result<OpenRouterUsage, SpendPanelError> {
262        let base_url = self.api_base(ctx);
263        let credits = Self::get_json::<OpenRouterCreditsResponse>(
264            client,
265            format!("{}/credits", base_url),
266            api_key,
267            ctx,
268        )
269        .await?;
270        let key_data = Self::fetch_key_data(client, &base_url, api_key, ctx).await;
271
272        Ok(OpenRouterUsage {
273            total_credits: credits.data.total_credits,
274            total_usage: credits.data.total_usage,
275            balance: credits.data.balance(),
276            key_data_fetched: key_data.is_some(),
277            key_limit: key_data.as_ref().and_then(|data| data.limit),
278            key_usage: key_data.as_ref().and_then(|data| data.usage),
279            key_usage_daily: key_data.as_ref().and_then(|data| data.usage_daily),
280            key_usage_weekly: key_data.as_ref().and_then(|data| data.usage_weekly),
281            key_usage_monthly: key_data.as_ref().and_then(|data| data.usage_monthly),
282            rate_limit: key_data.and_then(|data| data.rate_limit),
283        })
284    }
285
286    fn snapshot_from_usage(usage: OpenRouterUsage) -> UsageSnapshot {
287        let mut snapshot = UsageSnapshot::new("openrouter");
288
289        let mut credits = CreditsSnapshot::new(usage.balance, "USD");
290        credits.total = Some(usage.total_credits);
291        credits.used = Some(usage.total_usage);
292        snapshot.credits = Some(credits);
293
294        if let Some(ratio) = usage.key_used_ratio() {
295            snapshot.primary_rate_window = Some(RateWindow {
296                label: "API key limit".into(),
297                window_minutes: 0,
298                usage_ratio: ratio,
299                limit: None,
300                used: None,
301                remaining: None,
302                resets_at: None,
303                status: RateWindowStatus::from_ratio(ratio),
304            });
305            snapshot.cost = Some(CostSnapshot {
306                total_cost: usage.key_usage_monthly.or(usage.key_usage),
307                currency: "USD".into(),
308                daily_costs: Vec::new(),
309                spend_limit: Some(SpendLimit {
310                    limit: usage.key_limit.unwrap(),
311                    used: usage.key_usage.unwrap(),
312                    period: "api-key".into(),
313                }),
314            });
315        } else if usage.key_usage_monthly.is_some() || usage.key_usage.is_some() {
316            snapshot.cost = Some(CostSnapshot {
317                total_cost: usage.key_usage_monthly.or(usage.key_usage),
318                currency: "USD".into(),
319                daily_costs: Vec::new(),
320                spend_limit: None,
321            });
322        }
323
324        if let Some(rate_limit) = usage.rate_limit {
325            snapshot.plan = Some(crate::model::PlanInfo {
326                name: "OpenRouter API".into(),
327                tier: None,
328                features: vec![format!(
329                    "{} requests per {}",
330                    rate_limit.requests, rate_limit.interval
331                )],
332                price: None,
333                currency: None,
334                billing_period: None,
335            });
336        }
337
338        snapshot
339    }
340}
341
342impl Default for OpenRouterProvider {
343    fn default() -> Self {
344        Self::new()
345    }
346}
347
348#[async_trait]
349impl UsageProvider for OpenRouterProvider {
350    fn metadata(&self) -> &ProviderMetadata {
351        &self.metadata
352    }
353
354    fn detect_credentials(&self) -> bool {
355        Self::detect_credentials_from(std::env::var("OPENROUTER_API_KEY").ok().as_deref())
356    }
357
358    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
359        let api_key = Self::resolve_api_key(ctx)?;
360        let client = Self::build_client(ctx)?;
361        let usage = self.fetch_usage_data(&client, &api_key, ctx).await?;
362        Ok(Self::snapshot_from_usage(usage))
363    }
364}
365
366#[cfg(test)]
367mod tests {
368    use super::*;
369    use pretty_assertions::assert_eq;
370    use wiremock::matchers::{header, method, path};
371    use wiremock::{Mock, MockServer, ResponseTemplate};
372
373    #[test]
374    fn test_provider_metadata() {
375        let provider = OpenRouterProvider::new();
376        let meta = provider.metadata();
377        assert_eq!(meta.id, "openrouter");
378        assert_eq!(meta.name, "OpenRouter");
379        assert!(meta.auth_methods.contains(&"api_key"));
380    }
381
382    #[test]
383    fn test_clean_trims_and_unquotes() {
384        assert_eq!(OpenRouterProvider::clean("  sk-or  "), "sk-or");
385        assert_eq!(OpenRouterProvider::clean("\"sk-or\""), "sk-or");
386        assert_eq!(OpenRouterProvider::clean("'sk-or'"), "sk-or");
387    }
388
389    #[test]
390    fn test_resolve_api_key_from_context_token() {
391        let mut ctx = ProviderContext::new();
392        ctx.config.insert("token".into(), "sk-or-token".into());
393        assert_eq!(
394            OpenRouterProvider::resolve_api_key(&ctx).unwrap(),
395            "sk-or-token"
396        );
397    }
398
399    #[test]
400    fn test_api_base_accepts_bare_host() {
401        let mut ctx = ProviderContext::new();
402        ctx.config
403            .insert("api_url".into(), "openrouter.example/api/v1/".into());
404        assert_eq!(
405            OpenRouterProvider::new().api_base(&ctx),
406            "https://openrouter.example/api/v1"
407        );
408    }
409
410    #[test]
411    fn test_snapshot_with_key_quota() {
412        let snapshot = OpenRouterProvider::snapshot_from_usage(OpenRouterUsage {
413            total_credits: 100.0,
414            total_usage: 40.0,
415            balance: 60.0,
416            key_data_fetched: true,
417            key_limit: Some(20.0),
418            key_usage: Some(5.0),
419            key_usage_daily: Some(0.12),
420            key_usage_weekly: Some(0.74),
421            key_usage_monthly: Some(4.56),
422            rate_limit: Some(OpenRouterRateLimit {
423                requests: 120,
424                interval: "10s".into(),
425            }),
426        });
427        assert_eq!(snapshot.credits.as_ref().unwrap().balance, 60.0);
428        assert_eq!(snapshot.credits.as_ref().unwrap().used, Some(40.0));
429        assert_eq!(snapshot.primary_rate_window.unwrap().usage_ratio, 0.25);
430        assert_eq!(snapshot.cost.as_ref().unwrap().total_cost, Some(4.56));
431        assert_eq!(snapshot.cost.unwrap().spend_limit.unwrap().used, 5.0);
432        assert_eq!(snapshot.plan.unwrap().features[0], "120 requests per 10s");
433    }
434
435    #[test]
436    fn test_snapshot_without_key_quota_omits_primary_window() {
437        let snapshot = OpenRouterProvider::snapshot_from_usage(OpenRouterUsage {
438            total_credits: 50.0,
439            total_usage: 45.0,
440            balance: 5.0,
441            key_data_fetched: false,
442            key_limit: None,
443            key_usage: None,
444            key_usage_daily: None,
445            key_usage_weekly: None,
446            key_usage_monthly: None,
447            rate_limit: None,
448        });
449        assert!(snapshot.primary_rate_window.is_none());
450        assert_eq!(snapshot.credits.unwrap().balance, 5.0);
451    }
452
453    #[tokio::test]
454    async fn test_fetch_usage_success_with_key_enrichment() {
455        let server = MockServer::start().await;
456        Mock::given(method("GET"))
457            .and(path("/api/v1/credits"))
458            .and(header("authorization", "Bearer sk-or-test"))
459            .and(header("accept", "application/json"))
460            .and(header("http-referer", "https://usage.example"))
461            .and(header("x-title", "UsageMonitor QA"))
462            .respond_with(ResponseTemplate::new(200).set_body_raw(
463                r#"{"data":{"total_credits":100,"total_usage":40}}"#,
464                "application/json",
465            ))
466            .mount(&server)
467            .await;
468        Mock::given(method("GET"))
469            .and(path("/api/v1/key"))
470            .and(header("authorization", "Bearer sk-or-test"))
471            .respond_with(ResponseTemplate::new(200).set_body_raw(
472                r#"{"data":{"limit":20,"usage":0.5,"usage_daily":0.12,"usage_weekly":0.74,"usage_monthly":4.56,"rate_limit":{"requests":120,"interval":"10s"}}}"#,
473                "application/json",
474            ))
475            .mount(&server)
476            .await;
477
478        let provider = OpenRouterProvider::with_base_url(&format!("{}/api/v1", server.uri()));
479        let mut ctx = ProviderContext::with_api_key("sk-or-test");
480        ctx.config
481            .insert("http_referer".into(), " https://usage.example ".into());
482        ctx.config
483            .insert("x_title".into(), "UsageMonitor QA".into());
484        let snapshot = provider.fetch_usage(&ctx).await.unwrap();
485
486        assert_eq!(snapshot.credits.as_ref().unwrap().balance, 60.0);
487        assert_eq!(snapshot.primary_rate_window.unwrap().usage_ratio, 0.025);
488        assert_eq!(snapshot.cost.unwrap().total_cost, Some(4.56));
489    }
490
491    #[tokio::test]
492    async fn test_fetch_usage_key_failure_keeps_credits() {
493        let server = MockServer::start().await;
494        Mock::given(method("GET"))
495            .and(path("/api/v1/credits"))
496            .respond_with(ResponseTemplate::new(200).set_body_raw(
497                r#"{"data":{"total_credits":100,"total_usage":40}}"#,
498                "application/json",
499            ))
500            .mount(&server)
501            .await;
502        Mock::given(method("GET"))
503            .and(path("/api/v1/key"))
504            .respond_with(ResponseTemplate::new(500))
505            .mount(&server)
506            .await;
507
508        let provider = OpenRouterProvider::with_base_url(&format!("{}/api/v1", server.uri()));
509        let snapshot = provider
510            .fetch_usage(&ProviderContext::with_api_key("sk-or-test"))
511            .await
512            .unwrap();
513        assert_eq!(snapshot.credits.unwrap().balance, 60.0);
514        assert!(snapshot.primary_rate_window.is_none());
515    }
516
517    #[tokio::test]
518    async fn test_fetch_usage_401_is_auth_failed() {
519        let server = MockServer::start().await;
520        Mock::given(method("GET"))
521            .and(path("/api/v1/credits"))
522            .respond_with(ResponseTemplate::new(401))
523            .mount(&server)
524            .await;
525
526        let provider = OpenRouterProvider::with_base_url(&format!("{}/api/v1", server.uri()));
527        let err = provider
528            .fetch_usage(&ProviderContext::with_api_key("bad"))
529            .await
530            .unwrap_err();
531        assert!(matches!(err, SpendPanelError::AuthFailed(_, _)));
532    }
533}