usage_monitor_cli/provider/
ollama.rs

1//! Ollama cloud usage provider (ollama.com/settings, browser cookie auth).
2//!
3//! Ports CodexBar's scrape of the settings page: it reads the "Session usage"
4//! (or "Hourly usage") and "Weekly usage" percentages plus the plan name.
5
6use async_trait::async_trait;
7
8use crate::error::SpendPanelError;
9use crate::model::{PlanInfo, RateWindow, UsageSnapshot};
10use crate::provider::{ProviderContext, ProviderMetadata, UsageProvider};
11
12/// Finds the first `NN%` (allowing decimals) appearing after `label` in `html`.
13fn percent_after(html: &str, label: &str) -> Option<f64> {
14    let start = html.find(label)? + label.len();
15    // Bound the search window so a later block's percent isn't misattributed.
16    let end = (start + 800).min(html.len());
17    let window = &html[start..end];
18    let bytes = window.as_bytes();
19    let mut i = 0;
20    while i < bytes.len() {
21        if bytes[i] == b'%' {
22            // Walk back over an optional number immediately before '%'.
23            let mut j = i;
24            while j > 0 {
25                let c = bytes[j - 1];
26                if c.is_ascii_digit() || c == b'.' {
27                    j -= 1;
28                } else {
29                    break;
30                }
31            }
32            let parsed = (j < i).then(|| window[j..i].parse::<f64>().ok()).flatten();
33            if let Some(value) = parsed {
34                return Some(value.clamp(0.0, 100.0));
35            }
36        }
37        i += 1;
38    }
39    None
40}
41
42/// Extracts the inner text of the first `<span>…</span>` after `anchor`.
43fn span_after(html: &str, anchor: &str) -> Option<String> {
44    let start = html.find(anchor)? + anchor.len();
45    let rest = &html[start..];
46    let open = rest.find("<span")?;
47    let after_open = &rest[open..];
48    let gt = after_open.find('>')? + 1;
49    let inner = &after_open[gt..];
50    let close = inner.find('<')?;
51    let text = inner[..close].trim();
52    if text.is_empty() {
53        None
54    } else {
55        Some(text.to_string())
56    }
57}
58
59/// Ollama cloud usage provider.
60pub struct OllamaProvider {
61    metadata: ProviderMetadata,
62    base_url: Option<String>,
63}
64
65impl OllamaProvider {
66    pub fn new() -> Self {
67        Self {
68            metadata: ProviderMetadata {
69                id: "ollama",
70                name: "Ollama",
71                description: "Ollama cloud session/weekly usage monitor (browser cookie)",
72                auth_methods: &["cookie", "env"],
73                website: Some("https://ollama.com"),
74            },
75            base_url: None,
76        }
77    }
78
79    pub fn with_base_url(url: &str) -> Self {
80        let mut p = Self::new();
81        p.base_url = Some(url.to_string());
82        p
83    }
84
85    fn api_base(&self) -> &str {
86        self.base_url.as_deref().unwrap_or("https://ollama.com")
87    }
88
89    fn clean(raw: &str) -> String {
90        let mut v = raw.trim();
91        if v.len() >= 2
92            && ((v.starts_with('"') && v.ends_with('"'))
93                || (v.starts_with('\'') && v.ends_with('\'')))
94        {
95            v = &v[1..v.len() - 1];
96        }
97        v.trim().to_string()
98    }
99
100    fn resolve_cookie(ctx: &ProviderContext) -> Result<String, SpendPanelError> {
101        for key in ["cookie", "token"] {
102            if let Some(v) = ctx.config.get(key) {
103                let c = Self::clean(v);
104                if !c.is_empty() {
105                    return Ok(c);
106                }
107            }
108        }
109        if let Ok(v) = std::env::var("OLLAMA_COOKIE") {
110            let c = Self::clean(&v);
111            if !c.is_empty() {
112                return Ok(c);
113            }
114        }
115        Err(SpendPanelError::AuthFailed(
116            "ollama".into(),
117            "no session cookie in cookie config or OLLAMA_COOKIE".into(),
118        ))
119    }
120
121    fn build_client(ctx: &ProviderContext) -> Result<reqwest::Client, SpendPanelError> {
122        reqwest::Client::builder()
123            .timeout(std::time::Duration::from_secs(ctx.timeout_secs))
124            .build()
125            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))
126    }
127
128    fn parse(html: &str) -> Result<UsageSnapshot, SpendPanelError> {
129        let session =
130            percent_after(html, "Session usage").or_else(|| percent_after(html, "Hourly usage"));
131        let weekly = percent_after(html, "Weekly usage");
132
133        if session.is_none() && weekly.is_none() {
134            if html.contains("Sign in") || html.contains("sign in") {
135                return Err(SpendPanelError::AuthFailed(
136                    "ollama".into(),
137                    "not logged in to ollama.com (session cookie missing/expired)".into(),
138                ));
139            }
140            return Err(SpendPanelError::ParseError(
141                "ollama".into(),
142                "no usage data found on settings page".into(),
143            ));
144        }
145
146        let mut snapshot = UsageSnapshot::new("ollama");
147        if let Some(s) = session {
148            snapshot.primary_rate_window =
149                Some(RateWindow::new(s.round() as u64, 100, "Session", 5 * 60));
150        }
151        if let Some(w) = weekly {
152            snapshot.secondary_rate_window = Some(RateWindow::new(
153                w.round() as u64,
154                100,
155                "Weekly",
156                7 * 24 * 60,
157            ));
158        }
159        if let Some(plan) = span_after(html, "Cloud Usage") {
160            snapshot.plan = Some(PlanInfo {
161                name: plan,
162                tier: None,
163                features: Vec::new(),
164                price: None,
165                currency: None,
166                billing_period: None,
167            });
168        }
169        Ok(snapshot)
170    }
171}
172
173impl Default for OllamaProvider {
174    fn default() -> Self {
175        Self::new()
176    }
177}
178
179#[async_trait]
180impl UsageProvider for OllamaProvider {
181    fn metadata(&self) -> &ProviderMetadata {
182        &self.metadata
183    }
184
185    fn detect_credentials(&self) -> bool {
186        std::env::var("OLLAMA_COOKIE")
187            .map(|v| !v.trim().is_empty())
188            .unwrap_or(false)
189    }
190
191    async fn fetch_usage(&self, ctx: &ProviderContext) -> Result<UsageSnapshot, SpendPanelError> {
192        let cookie = Self::resolve_cookie(ctx)?;
193        let client = Self::build_client(ctx)?;
194        let url = format!("{}/settings", self.api_base().trim_end_matches('/'));
195        let resp = client
196            .get(url)
197            .header("Cookie", cookie)
198            .header("Accept", "text/html")
199            .send()
200            .await
201            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
202        let status = resp.status();
203        let body = resp
204            .text()
205            .await
206            .map_err(|e| SpendPanelError::NetworkError(e.to_string()))?;
207        if status == reqwest::StatusCode::UNAUTHORIZED || status == reqwest::StatusCode::FORBIDDEN {
208            return Err(SpendPanelError::AuthFailed(
209                "ollama".into(),
210                format!("session cookie rejected (HTTP {})", status.as_u16()),
211            ));
212        }
213        if !status.is_success() {
214            return Err(SpendPanelError::ProviderError(
215                "ollama".into(),
216                format!("HTTP {}: {}", status, body),
217            ));
218        }
219        Self::parse(&body)
220    }
221}
222
223#[cfg(test)]
224mod tests {
225    use super::*;
226    use pretty_assertions::assert_eq;
227    use wiremock::matchers::{method, path};
228    use wiremock::{Mock, MockServer, ResponseTemplate};
229
230    const HTML: &str = r#"<html><body>
231      <div>Cloud Usage <span class="x">Pro</span></div>
232      <div>Session usage <div class="bar">42%</div> resets soon</div>
233      <div>Weekly usage <div class="bar">7.5%</div></div>
234    </body></html>"#;
235
236    #[test]
237    fn test_metadata() {
238        assert_eq!(OllamaProvider::new().metadata().id, "ollama");
239    }
240
241    #[test]
242    fn test_percent_after() {
243        assert_eq!(percent_after(HTML, "Session usage"), Some(42.0));
244        assert_eq!(percent_after(HTML, "Weekly usage"), Some(7.5));
245    }
246
247    #[test]
248    fn test_parse() {
249        let snap = OllamaProvider::parse(HTML).unwrap();
250        assert_eq!(snap.primary_rate_window.unwrap().used, Some(42));
251        assert_eq!(snap.secondary_rate_window.unwrap().used, Some(8));
252        assert_eq!(snap.plan.unwrap().name, "Pro");
253    }
254
255    #[test]
256    fn test_parse_signed_out() {
257        let err = OllamaProvider::parse("<html>Please Sign in</html>").unwrap_err();
258        assert!(matches!(err, SpendPanelError::AuthFailed(_, _)));
259    }
260
261    #[tokio::test]
262    async fn test_fetch_success() {
263        let server = MockServer::start().await;
264        Mock::given(method("GET"))
265            .and(path("/settings"))
266            .respond_with(ResponseTemplate::new(200).set_body_raw(HTML, "text/html"))
267            .mount(&server)
268            .await;
269        let provider = OllamaProvider::with_base_url(&server.uri());
270        let mut ctx = ProviderContext::new();
271        ctx.config.insert("cookie".into(), "sid=abc".into());
272        let snap = provider.fetch_usage(&ctx).await.unwrap();
273        assert_eq!(snap.primary_rate_window.unwrap().used, Some(42));
274    }
275
276    #[tokio::test]
277    async fn test_fetch_401() {
278        let server = MockServer::start().await;
279        Mock::given(method("GET"))
280            .and(path("/settings"))
281            .respond_with(ResponseTemplate::new(401))
282            .mount(&server)
283            .await;
284        let provider = OllamaProvider::with_base_url(&server.uri());
285        let mut ctx = ProviderContext::new();
286        ctx.config.insert("cookie".into(), "bad".into());
287        assert!(matches!(
288            provider.fetch_usage(&ctx).await.unwrap_err(),
289            SpendPanelError::AuthFailed(_, _)
290        ));
291    }
292}
usage_monitor_cli/provider/ollama.rs

usage_monitor_cli/provider/
ollama.rs
