Skip to main content

ai_usagebar/nous/
fetch.rs

1//! Nous account transport and response classification.
2
3use chrono::{DateTime, Utc};
4use thiserror::Error;
5
6use super::credentials::CredentialStore;
7use super::oauth;
8use super::types::{AccountSnapshot, parse_account};
9
10pub const ACCOUNT_URL: &str = "https://portal.nousresearch.com/api/oauth/account";
11
12#[derive(Debug, Clone, PartialEq, Eq)]
13pub struct Endpoints {
14    pub account: String,
15    pub token: String,
16}
17
18impl Default for Endpoints {
19    fn default() -> Self {
20        Self {
21            account: ACCOUNT_URL.into(),
22            token: oauth::TOKEN_URL.into(),
23        }
24    }
25}
26
27#[derive(Debug, Error, Clone, PartialEq, Eq)]
28pub enum FetchError {
29    #[error("Nous authentication failed")]
30    Authentication,
31    #[error("Nous account endpoint rate limited the request")]
32    RateLimited,
33    #[error("Nous account endpoint is temporarily unavailable")]
34    Transient,
35    #[error("Nous account response schema mismatch")]
36    Schema,
37    #[error("Nous account HTTP status {0}")]
38    HttpStatus(u16),
39    #[error("Nous account network transport failed")]
40    Transport,
41    #[error("Nous account response exceeded the body limit")]
42    BodyLimit,
43}
44
45impl From<FetchError> for crate::error::AppError {
46    fn from(error: FetchError) -> Self {
47        use crate::error::{AUTH_FAILURE_MESSAGE, AppError};
48        match error {
49            FetchError::Authentication => AppError::Credentials(AUTH_FAILURE_MESSAGE.to_string()),
50            FetchError::RateLimited => AppError::Http {
51                status: 429,
52                body: "Nous Research request was rate limited".into(),
53            },
54            FetchError::Transient | FetchError::Transport => {
55                AppError::Transport("Nous Research request failed".into())
56            }
57            FetchError::Schema => {
58                AppError::Schema("Nous Research account response schema mismatch".into())
59            }
60            FetchError::HttpStatus(status) => AppError::Http {
61                status,
62                body: "Nous Research request failed".into(),
63            },
64            FetchError::BodyLimit => {
65                AppError::Schema("Nous Research response exceeded the body limit".into())
66            }
67        }
68    }
69}
70
71pub async fn fetch_account(
72    client: &reqwest::Client,
73    access_token: &str,
74    endpoints: &Endpoints,
75) -> Result<AccountSnapshot, FetchError> {
76    if access_token.trim().is_empty() {
77        return Err(FetchError::Authentication);
78    }
79    let response = client
80        .get(&endpoints.account)
81        .bearer_auth(access_token)
82        .header("accept", "application/json")
83        .send()
84        .await
85        .map_err(|_| FetchError::Transport)?;
86    let status = response.status();
87    let body = crate::vendor::read_body_capped(response, crate::vendor::MAX_BODY_BYTES)
88        .await
89        .map_err(|error| {
90            if error.to_string().contains("exceeds") {
91                FetchError::BodyLimit
92            } else {
93                FetchError::Transport
94            }
95        })?;
96    if !status.is_success() {
97        return Err(classify_status(status.as_u16()));
98    }
99    let value: serde_json::Value = serde_json::from_slice(&body).map_err(|_| FetchError::Schema)?;
100    parse_account(&value).map_err(|_| FetchError::Schema)
101}
102
103pub async fn fetch_account_with_refresh(
104    client: &reqwest::Client,
105    store: &CredentialStore,
106    endpoints: &Endpoints,
107    now: DateTime<Utc>,
108) -> Result<AccountSnapshot, FetchError> {
109    let credential = oauth::refresh_if_needed(client, store, &endpoints.token, now)
110        .await
111        .map_err(map_oauth_error)?;
112    fetch_account(client, &credential.access_token, endpoints).await
113}
114
115fn classify_status(status: u16) -> FetchError {
116    match status {
117        401 | 403 => FetchError::Authentication,
118        429 => FetchError::RateLimited,
119        500..=599 => FetchError::Transient,
120        other => FetchError::HttpStatus(other),
121    }
122}
123
124fn map_oauth_error(error: oauth::OAuthError) -> FetchError {
125    match error {
126        oauth::OAuthError::Transport => FetchError::Transport,
127        oauth::OAuthError::RefreshTokenRejected
128        | oauth::OAuthError::Credentials
129        | oauth::OAuthError::AccessDenied
130        | oauth::OAuthError::ExpiredToken => FetchError::Authentication,
131        oauth::OAuthError::Schema => FetchError::Schema,
132        oauth::OAuthError::HttpStatus(429) => FetchError::RateLimited,
133        oauth::OAuthError::HttpStatus(status) if status >= 500 => FetchError::Transient,
134        oauth::OAuthError::HttpStatus(status) => FetchError::HttpStatus(status),
135        oauth::OAuthError::UnknownOAuthError | oauth::OAuthError::Deadline => FetchError::Schema,
136    }
137}
138
139#[cfg(test)]
140mod tests {
141    use chrono::{Duration as ChronoDuration, Utc};
142    use tempfile::TempDir;
143
144    use super::*;
145    use crate::nous::credentials::{CredentialDocument, CredentialStore, NousCredential};
146
147    #[test]
148    fn app_error_projection_preserves_non_auth_failure_classes() {
149        assert!(matches!(
150            crate::error::AppError::from(FetchError::RateLimited),
151            crate::error::AppError::Http { status: 429, .. }
152        ));
153        assert!(matches!(
154            crate::error::AppError::from(FetchError::Transport),
155            crate::error::AppError::Transport(_)
156        ));
157        assert!(matches!(
158            crate::error::AppError::from(FetchError::Schema),
159            crate::error::AppError::Schema(_)
160        ));
161    }
162
163    #[tokio::test]
164    async fn account_request_uses_exact_path_bearer_and_accept_headers() {
165        let mut server = mockito::Server::new_async().await;
166        let mock = server
167            .mock("GET", "/api/oauth/account")
168            .match_header("authorization", "Bearer test-access-token")
169            .match_header("accept", "application/json")
170            .with_status(200)
171            .with_body(include_str!("../../tests/fixtures/nous/account.json"))
172            .create_async()
173            .await;
174        let endpoints = Endpoints {
175            account: format!("{}/api/oauth/account", server.url()),
176            token: format!("{}/token", server.url()),
177        };
178
179        let snapshot = fetch_account(&reqwest::Client::new(), "test-access-token", &endpoints)
180            .await
181            .unwrap();
182        assert_eq!(snapshot.plan.as_deref(), Some("Pro"));
183        mock.assert_async().await;
184    }
185
186    #[tokio::test]
187    async fn account_statuses_are_classified_without_retaining_response_bodies() {
188        for (status, expected) in [
189            (401, FetchError::Authentication),
190            (403, FetchError::Authentication),
191            (429, FetchError::RateLimited),
192            (500, FetchError::Transient),
193        ] {
194            let mut server = mockito::Server::new_async().await;
195            server
196                .mock("GET", "/account")
197                .with_status(status)
198                .with_body("test-secret-response-body")
199                .create_async()
200                .await;
201            let endpoints = Endpoints {
202                account: format!("{}/account", server.url()),
203                token: format!("{}/token", server.url()),
204            };
205            let error = fetch_account(&reqwest::Client::new(), "test-access-token", &endpoints)
206                .await
207                .unwrap_err();
208            assert_eq!(error, expected);
209            assert!(!format!("{error:?}").contains("test-secret-response-body"));
210        }
211    }
212
213    #[tokio::test]
214    async fn malformed_success_is_a_schema_error_and_network_failure_is_transient() {
215        let mut server = mockito::Server::new_async().await;
216        server
217            .mock("GET", "/account")
218            .with_status(200)
219            .with_body(r#"{"error":"test-secret"}"#)
220            .create_async()
221            .await;
222        let endpoints = Endpoints {
223            account: format!("{}/account", server.url()),
224            token: format!("{}/token", server.url()),
225        };
226        assert_eq!(
227            fetch_account(&reqwest::Client::new(), "test-access-token", &endpoints)
228                .await
229                .unwrap_err(),
230            FetchError::Schema
231        );
232
233        let network = Endpoints {
234            account: "http://127.0.0.1:1/account".into(),
235            token: "http://127.0.0.1:1/token".into(),
236        };
237        assert_eq!(
238            fetch_account(&reqwest::Client::new(), "test-access-token", &network)
239                .await
240                .unwrap_err(),
241            FetchError::Transport
242        );
243    }
244
245    #[tokio::test]
246    async fn refresh_is_persisted_before_account_probe() {
247        let mut server = mockito::Server::new_async().await;
248        server
249            .mock("POST", "/token")
250            .match_header("x-nous-refresh-token", "test-old-refresh")
251            .with_status(200)
252            .with_body(r#"{"access_token":"test-new-access","refresh_token":"test-new-refresh","token_type":"Bearer","expires_in":3600}"#)
253            .create_async()
254            .await;
255        let account_mock = server
256            .mock("GET", "/account")
257            .match_header("authorization", "Bearer test-new-access")
258            .with_status(200)
259            .with_body(include_str!("../../tests/fixtures/nous/account.json"))
260            .create_async()
261            .await;
262        let root = TempDir::new().unwrap();
263        let path = root.path().join("config").join("credentials.json");
264        let store = CredentialStore::at(&path);
265        store
266            .write(&CredentialDocument::new(Some(NousCredential {
267                client_id: "hermes-cli".into(),
268                access_token: "test-old-access".into(),
269                refresh_token: "test-old-refresh".into(),
270                expires_at: Utc::now() + ChronoDuration::seconds(100),
271            })))
272            .unwrap();
273        let endpoints = Endpoints {
274            account: format!("{}/account", server.url()),
275            token: format!("{}/token", server.url()),
276        };
277
278        let snapshot =
279            fetch_account_with_refresh(&reqwest::Client::new(), &store, &endpoints, Utc::now())
280                .await
281                .unwrap();
282        assert_eq!(snapshot.plan.as_deref(), Some("Pro"));
283        assert_eq!(
284            store.read().unwrap().unwrap().nous.unwrap().refresh_token,
285            "test-new-refresh"
286        );
287        account_mock.assert_async().await;
288    }
289}