Skip to main content

codex_cost/
provider.rs

1use std::{
2    env,
3    path::{Path, PathBuf},
4};
5
6use walkdir::WalkDir;
7
8use crate::{
9    models::{UsageLoad, UsageWarning},
10    parser::parse_session_file,
11};
12
13pub trait UsageProvider: Send + Sync {
14    fn load_events(&self) -> UsageLoad;
15}
16
17#[derive(Debug, Clone)]
18pub struct LocalCodexProvider {
19    codex_home: PathBuf,
20}
21
22impl LocalCodexProvider {
23    pub fn new(codex_home: Option<PathBuf>) -> Self {
24        Self {
25            codex_home: codex_home.unwrap_or_else(default_codex_home),
26        }
27    }
28
29    pub fn codex_home(&self) -> &Path {
30        &self.codex_home
31    }
32}
33
34impl UsageProvider for LocalCodexProvider {
35    fn load_events(&self) -> UsageLoad {
36        if !self.codex_home.exists() {
37            return UsageLoad {
38                warnings: vec![UsageWarning {
39                    file: self.codex_home.display().to_string(),
40                    line: None,
41                    message: "Codex home not found".to_string(),
42                }],
43                ..UsageLoad::default()
44            };
45        }
46
47        let mut load = UsageLoad::default();
48        for path in discover_session_files(&self.codex_home) {
49            load.session_count += 1;
50            let parsed = parse_session_file(&path);
51            load.events.extend(parsed.events);
52            load.warnings.extend(parsed.warnings);
53        }
54
55        load
56    }
57}
58
59fn discover_session_files(codex_home: &Path) -> Vec<PathBuf> {
60    let mut files = Vec::new();
61    collect_jsonl_files(&codex_home.join("sessions"), &mut files);
62    collect_jsonl_files(&codex_home.join("archived_sessions"), &mut files);
63    files.sort();
64    files
65}
66
67fn collect_jsonl_files(root: &Path, files: &mut Vec<PathBuf>) {
68    if !root.exists() {
69        return;
70    }
71
72    for entry in WalkDir::new(root).into_iter().filter_map(Result::ok) {
73        if entry.file_type().is_file()
74            && entry
75                .path()
76                .extension()
77                .is_some_and(|extension| extension.eq_ignore_ascii_case("jsonl"))
78        {
79            files.push(entry.path().to_path_buf());
80        }
81    }
82}
83
84fn default_codex_home() -> PathBuf {
85    if let Some(home) = env::var_os("USERPROFILE").or_else(|| env::var_os("HOME")) {
86        return PathBuf::from(home).join(".codex");
87    }
88
89    PathBuf::from(".codex")
90}
91
92#[cfg(test)]
93mod tests {
94    use std::{fs, path::Path};
95
96    use tempfile::tempdir;
97
98    use super::{LocalCodexProvider, UsageProvider};
99
100    fn write_session(path: &Path, session_id: &str, cwd: &str, tokens: u64) {
101        if let Some(parent) = path.parent() {
102            fs::create_dir_all(parent).expect("create fixture directory");
103        }
104
105        let session_meta = serde_json::json!({
106            "timestamp": "2026-07-06T08:00:00.000Z",
107            "type": "session_meta",
108            "payload": {
109                "id": session_id,
110                "cwd": cwd,
111                "thread_name": format!("Thread {session_id}"),
112            },
113        });
114        let token_count = serde_json::json!({
115            "timestamp": "2026-07-06T08:01:00.000Z",
116            "type": "event_msg",
117            "payload": {
118                "type": "token_count",
119                "info": {
120                    "last_token_usage": {
121                        "total_tokens": tokens,
122                    },
123                },
124            },
125        });
126        fs::write(path, format!("{session_meta}\n{token_count}\n")).expect("write fixture session");
127    }
128
129    #[test]
130    fn discovers_active_and_archived_session_files() {
131        let root = tempdir().expect("create codex home");
132        write_session(
133            &root.path().join("sessions/2026/07/06/rollout-active.jsonl"),
134            "active",
135            "C:\\work\\active",
136            11,
137        );
138        write_session(
139            &root.path().join("archived_sessions/rollout-archived.jsonl"),
140            "archived",
141            "C:\\work\\archived",
142            17,
143        );
144
145        let provider = LocalCodexProvider::new(Some(root.path().to_path_buf()));
146        let loaded = provider.load_events();
147
148        assert_eq!(loaded.events.len(), 2);
149        assert_eq!(loaded.session_count, 2);
150        assert!(loaded.warnings.is_empty());
151        assert!(
152            loaded
153                .events
154                .iter()
155                .any(|event| event.project == "C:\\work\\active")
156        );
157        assert!(
158            loaded
159                .events
160                .iter()
161                .any(|event| event.project == "C:\\work\\archived")
162        );
163        assert_eq!(
164            loaded
165                .events
166                .iter()
167                .map(|event| event.usage.total_tokens)
168                .sum::<u64>(),
169            28
170        );
171    }
172
173    #[test]
174    fn missing_codex_home_returns_empty_load_with_warning() {
175        let root = tempdir().expect("create temp root");
176        let missing = root.path().join("missing-codex-home");
177
178        let provider = LocalCodexProvider::new(Some(missing));
179        let loaded = provider.load_events();
180
181        assert!(loaded.events.is_empty());
182        assert_eq!(loaded.session_count, 0);
183        assert_eq!(loaded.warnings.len(), 1);
184        assert!(loaded.warnings[0].message.contains("Codex home not found"));
185    }
186}