Skip to main content

codex_ops/
account_history.rs

1use crate::error::AppError;
2use crate::storage::write_sensitive_file;
3use chrono::{DateTime, SecondsFormat, Utc};
4use serde::{Deserialize, Serialize};
5use std::fs;
6use std::io;
7use std::path::Path;
8
9pub const ACCOUNT_HISTORY_STORE_VERSION: u8 = 1;
10pub const DEFAULT_ACCOUNT_SOURCE: &str = "auth.json";
11pub const AUTH_SELECT_SOURCE: &str = "auth select";
12
13#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
14#[serde(rename_all = "camelCase")]
15pub struct AccountHistoryAccount {
16    pub account_id: String,
17    pub observed_at: String,
18    pub source: String,
19    #[serde(skip_serializing_if = "Option::is_none")]
20    pub name: Option<String>,
21    #[serde(skip_serializing_if = "Option::is_none")]
22    pub email: Option<String>,
23    #[serde(skip_serializing_if = "Option::is_none")]
24    pub plan_type: Option<String>,
25}
26
27impl AccountHistoryAccount {
28    pub fn auth_json(
29        account_id: String,
30        observed_at: DateTime<Utc>,
31        name: Option<String>,
32        email: Option<String>,
33        plan_type: Option<String>,
34    ) -> Self {
35        Self {
36            account_id,
37            observed_at: format_account_history_iso(observed_at),
38            source: DEFAULT_ACCOUNT_SOURCE.to_string(),
39            name,
40            email,
41            plan_type,
42        }
43    }
44}
45
46#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
47#[serde(rename_all = "camelCase")]
48pub struct AccountHistorySwitchEvent {
49    pub timestamp: String,
50    pub from_account_id: String,
51    pub to_account_id: String,
52    pub source: String,
53}
54
55#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
56#[serde(rename_all = "camelCase")]
57pub struct AccountHistoryStore {
58    pub version: u8,
59    #[serde(skip_serializing_if = "Option::is_none")]
60    pub default_account: Option<AccountHistoryAccount>,
61    pub switches: Vec<AccountHistorySwitchEvent>,
62}
63
64#[derive(Debug, Clone)]
65pub struct UsageAccountHistory {
66    default_account_id: Option<String>,
67    switches: Vec<UsageAccountSwitch>,
68}
69
70impl UsageAccountHistory {
71    pub fn account_id_at(&self, timestamp: DateTime<Utc>) -> Option<String> {
72        let mut account_id = self.default_account_id.clone();
73        for entry in &self.switches {
74            if entry.timestamp > timestamp {
75                break;
76            }
77            account_id = Some(entry.to_account_id.clone());
78        }
79        account_id
80    }
81}
82
83#[derive(Debug, Clone)]
84struct UsageAccountSwitch {
85    timestamp: DateTime<Utc>,
86    to_account_id: String,
87}
88
89pub fn read_account_history_store(path: &Path) -> Result<AccountHistoryStore, AppError> {
90    let content = match fs::read_to_string(path) {
91        Ok(content) => content,
92        Err(error) if error.kind() == io::ErrorKind::NotFound => {
93            return Ok(empty_account_history_store());
94        }
95        Err(error) => return Err(AppError::new(error.to_string())),
96    };
97
98    parse_account_history_store(&content, path)
99}
100
101pub fn write_account_history_store(
102    path: &Path,
103    store: &AccountHistoryStore,
104) -> Result<(), AppError> {
105    let normalized = normalize_account_history_store(store.clone())?;
106    let content = serde_json::to_string_pretty(&normalized)
107        .map_err(|error| AppError::new(error.to_string()))?;
108    write_sensitive_file(path, &format!("{content}\n"))
109        .map_err(|error| AppError::new(error.to_string()))
110}
111
112pub fn ensure_default_account(
113    store: AccountHistoryStore,
114    account: AccountHistoryAccount,
115) -> Result<AccountHistoryStore, AppError> {
116    let mut normalized = normalize_account_history_store(store)?;
117    if normalized.default_account.is_none() {
118        normalized.default_account = Some(account);
119    }
120    normalize_account_history_store(normalized)
121}
122
123pub fn ensure_default_account_in_file(
124    path: &Path,
125    account: AccountHistoryAccount,
126) -> Result<AccountHistoryStore, AppError> {
127    let current = read_account_history_store(path)?;
128    let ensured = ensure_default_account(current, account)?;
129    write_account_history_store(path, &ensured)?;
130    Ok(ensured)
131}
132
133pub fn record_auth_select_switch(
134    store: AccountHistoryStore,
135    default_account: AccountHistoryAccount,
136    from_account_id: &str,
137    to_account_id: &str,
138    timestamp: DateTime<Utc>,
139) -> Result<AccountHistoryStore, AppError> {
140    let mut ensured = ensure_default_account(store, default_account)?;
141    ensured.switches.push(AccountHistorySwitchEvent {
142        timestamp: format_account_history_iso(timestamp),
143        from_account_id: from_account_id.to_string(),
144        to_account_id: to_account_id.to_string(),
145        source: AUTH_SELECT_SOURCE.to_string(),
146    });
147    normalize_account_history_store(ensured)
148}
149
150pub fn read_optional_usage_account_history(
151    path: &Path,
152) -> Result<Option<UsageAccountHistory>, AppError> {
153    if !path.exists() {
154        return Ok(None);
155    }
156
157    usage_account_history_from_store(read_account_history_store(path)?)
158}
159
160pub fn usage_account_history_from_store(
161    store: AccountHistoryStore,
162) -> Result<Option<UsageAccountHistory>, AppError> {
163    let store = normalize_account_history_store(store)?;
164    if store.default_account.is_none() && store.switches.is_empty() {
165        return Ok(None);
166    }
167
168    let default_account_id = store.default_account.map(|account| account.account_id);
169    let switches = store
170        .switches
171        .into_iter()
172        .map(|entry| {
173            Ok(UsageAccountSwitch {
174                timestamp: parse_timestamp(&entry.timestamp, "switch.timestamp")?,
175                to_account_id: entry.to_account_id,
176            })
177        })
178        .collect::<Result<Vec<_>, AppError>>()?;
179
180    Ok(Some(UsageAccountHistory {
181        default_account_id,
182        switches,
183    }))
184}
185
186pub fn read_default_account_label_for(
187    path: &Path,
188    account_id: &str,
189) -> Result<Option<String>, AppError> {
190    let store = read_account_history_store(path)?;
191    let Some(account) = store.default_account else {
192        return Ok(None);
193    };
194    if account.account_id != account_id {
195        return Ok(None);
196    }
197    let label = account.email.or(account.name);
198    Ok(label.map(|label| format!("{label}({account_id})")))
199}
200
201pub fn format_account_history_iso(date: DateTime<Utc>) -> String {
202    date.to_rfc3339_opts(SecondsFormat::Millis, true)
203}
204
205fn empty_account_history_store() -> AccountHistoryStore {
206    AccountHistoryStore {
207        version: ACCOUNT_HISTORY_STORE_VERSION,
208        default_account: None,
209        switches: Vec::new(),
210    }
211}
212
213fn parse_account_history_store(
214    content: &str,
215    path: &Path,
216) -> Result<AccountHistoryStore, AppError> {
217    if content.trim().is_empty() {
218        return Ok(empty_account_history_store());
219    }
220
221    let parsed: AccountHistoryStore = serde_json::from_str(content).map_err(|error| {
222        AppError::new(format!(
223            "Failed to parse {}: {}",
224            path_to_string(path),
225            error
226        ))
227    })?;
228
229    if parsed.version != ACCOUNT_HISTORY_STORE_VERSION {
230        return Err(AppError::new(format!(
231            "Unsupported auth account history version in {}: {}.",
232            path_to_string(path),
233            parsed.version
234        )));
235    }
236
237    normalize_account_history_store(parsed)
238}
239
240fn normalize_account_history_store(
241    mut store: AccountHistoryStore,
242) -> Result<AccountHistoryStore, AppError> {
243    if let Some(default_account) = &mut store.default_account {
244        default_account.account_id =
245            normalize_required_account_id(&default_account.account_id, "default account id")?;
246        if default_account.source != DEFAULT_ACCOUNT_SOURCE {
247            return Err(AppError::new(
248                "Expected defaultAccount.source to be auth.json.",
249            ));
250        }
251        parse_timestamp(&default_account.observed_at, "defaultAccount.observedAt")?;
252        default_account.name = normalize_optional_string(default_account.name.take());
253        default_account.email = normalize_optional_string(default_account.email.take());
254        default_account.plan_type = normalize_optional_string(default_account.plan_type.take());
255    }
256
257    for entry in &mut store.switches {
258        parse_timestamp(&entry.timestamp, "switch.timestamp")?;
259        entry.from_account_id =
260            normalize_required_account_id(&entry.from_account_id, "switch from account id")?;
261        entry.to_account_id =
262            normalize_required_account_id(&entry.to_account_id, "switch to account id")?;
263        if entry.source != AUTH_SELECT_SOURCE {
264            return Err(AppError::new("Expected switch.source to be auth select."));
265        }
266    }
267
268    store.switches.sort_by(|left, right| {
269        parse_timestamp(&left.timestamp, "switch.timestamp")
270            .expect("switch timestamp validated")
271            .cmp(
272                &parse_timestamp(&right.timestamp, "switch.timestamp")
273                    .expect("switch timestamp validated"),
274            )
275            .then_with(|| left.to_account_id.cmp(&right.to_account_id))
276            .then_with(|| left.from_account_id.cmp(&right.from_account_id))
277    });
278
279    Ok(AccountHistoryStore {
280        version: ACCOUNT_HISTORY_STORE_VERSION,
281        default_account: store.default_account,
282        switches: store.switches,
283    })
284}
285
286fn parse_timestamp(value: &str, path: &str) -> Result<DateTime<Utc>, AppError> {
287    DateTime::parse_from_rfc3339(value)
288        .map(|date| date.with_timezone(&Utc))
289        .map_err(|_| AppError::new(format!("Expected {path} to be a valid date string.")))
290}
291
292fn normalize_required_account_id(value: &str, label: &str) -> Result<String, AppError> {
293    let normalized = value.trim();
294    if normalized.is_empty() {
295        return Err(AppError::new(format!("{label} cannot be empty.")));
296    }
297    Ok(normalized.to_string())
298}
299
300fn normalize_optional_string(value: Option<String>) -> Option<String> {
301    value.and_then(|value| {
302        let trimmed = value.trim().to_string();
303        (!trimmed.is_empty()).then_some(trimmed)
304    })
305}
306
307fn path_to_string(path: &Path) -> String {
308    path.to_string_lossy().to_string()
309}
310
311#[cfg(test)]
312mod tests {
313    use super::*;
314    use std::path::PathBuf;
315    use std::time::{SystemTime, UNIX_EPOCH};
316
317    #[test]
318    fn missing_and_empty_history_files_read_as_empty_store() {
319        let root = temp_dir("codex-ops-account-history-empty");
320        fs::create_dir_all(&root).expect("create root");
321        let missing = root.join("missing.json");
322        let empty = root.join("empty.json");
323        fs::write(&empty, "  \n").expect("write empty");
324
325        assert_eq!(
326            read_account_history_store(&missing).expect("missing store"),
327            empty_account_history_store()
328        );
329        assert_eq!(
330            read_account_history_store(&empty).expect("empty store"),
331            empty_account_history_store()
332        );
333
334        let _ = fs::remove_dir_all(root);
335    }
336
337    #[test]
338    fn unsupported_version_is_rejected() {
339        let root = temp_dir("codex-ops-account-history-version");
340        fs::create_dir_all(&root).expect("create root");
341        let path = root.join("history.json");
342        fs::write(&path, r#"{"version":2,"switches":[]}"#).expect("write history");
343
344        let error = read_account_history_store(&path).expect_err("unsupported version");
345
346        assert!(error
347            .message()
348            .contains("Unsupported auth account history version"));
349        let _ = fs::remove_dir_all(root);
350    }
351
352    #[test]
353    fn switches_are_sorted_for_usage_attribution() {
354        let store = normalize_account_history_store(AccountHistoryStore {
355            version: ACCOUNT_HISTORY_STORE_VERSION,
356            default_account: Some(account("account-a", "2026-05-01T00:00:00.000Z")),
357            switches: vec![
358                switch("2026-05-03T00:00:00.000Z", "account-b", "account-c"),
359                switch("2026-05-02T00:00:00.000Z", "account-a", "account-b"),
360            ],
361        })
362        .expect("normalize");
363        let usage = usage_account_history_from_store(store)
364            .expect("usage history")
365            .expect("non-empty usage history");
366
367        assert_eq!(
368            usage.account_id_at(parse("2026-05-01T12:00:00.000Z")),
369            Some("account-a".to_string())
370        );
371        assert_eq!(
372            usage.account_id_at(parse("2026-05-02T12:00:00.000Z")),
373            Some("account-b".to_string())
374        );
375        assert_eq!(
376            usage.account_id_at(parse("2026-05-03T12:00:00.000Z")),
377            Some("account-c".to_string())
378        );
379    }
380
381    #[test]
382    fn ensure_default_account_initializes_and_writes_store() {
383        let root = temp_dir("codex-ops-account-history-ensure");
384        fs::create_dir_all(&root).expect("create root");
385        let path = root.join("history.json");
386
387        let store =
388            ensure_default_account_in_file(&path, account("account-a", "2026-05-01T00:00:00.000Z"))
389                .expect("ensure default");
390
391        assert_eq!(
392            store.default_account.expect("default").account_id,
393            "account-a"
394        );
395        let content = fs::read_to_string(&path).expect("read written history");
396        assert!(content.contains(r#""defaultAccount""#));
397        assert!(content.ends_with('\n'));
398
399        let _ = fs::remove_dir_all(root);
400    }
401
402    fn account(account_id: &str, observed_at: &str) -> AccountHistoryAccount {
403        AccountHistoryAccount {
404            account_id: account_id.to_string(),
405            observed_at: observed_at.to_string(),
406            source: DEFAULT_ACCOUNT_SOURCE.to_string(),
407            name: None,
408            email: None,
409            plan_type: None,
410        }
411    }
412
413    fn switch(timestamp: &str, from: &str, to: &str) -> AccountHistorySwitchEvent {
414        AccountHistorySwitchEvent {
415            timestamp: timestamp.to_string(),
416            from_account_id: from.to_string(),
417            to_account_id: to.to_string(),
418            source: AUTH_SELECT_SOURCE.to_string(),
419        }
420    }
421
422    fn parse(value: &str) -> DateTime<Utc> {
423        DateTime::parse_from_rfc3339(value)
424            .expect("timestamp")
425            .with_timezone(&Utc)
426    }
427
428    fn temp_dir(prefix: &str) -> PathBuf {
429        let millis = SystemTime::now()
430            .duration_since(UNIX_EPOCH)
431            .expect("system time")
432            .as_millis();
433        std::env::temp_dir().join(format!("{prefix}-{millis}-{}", std::process::id()))
434    }
435}