Skip to main content

warden/
config.rs

1//! `~/.warden/config.toml`. Every field is optional and has a default.
2//!
3//! Pricing lives here and only here — the binary never hardcodes a rate, and an
4//! unpriced model yields `None` rather than a misleading `0.0`.
5
6use std::collections::BTreeMap;
7use std::fmt;
8use std::io;
9use std::path::{Path, PathBuf};
10
11use serde::Deserialize;
12
13/// Loaded configuration, with defaults already applied.
14#[derive(Debug, Clone, Default, Deserialize)]
15#[serde(default)]
16pub struct Config {
17    pub general: General,
18    /// Source adapters keyed by adapter name, e.g. `claude-code`.
19    pub sources: BTreeMap<String, Source>,
20    /// Per-provider price tables keyed by provider, e.g. `anthropic`.
21    pub pricing: BTreeMap<String, BTreeMap<String, ModelPrice>>,
22}
23
24#[derive(Debug, Clone, Deserialize)]
25#[serde(default)]
26pub struct General {
27    /// Store location. `None` means "use the built-in default" (`~/.warden`).
28    pub data_dir: Option<PathBuf>,
29    /// Store prompt text alongside `text_hash`.
30    pub index_prompt_text: bool,
31}
32
33impl Default for General {
34    fn default() -> Self {
35        Self {
36            data_dir: None,
37            index_prompt_text: true,
38        }
39    }
40}
41
42#[derive(Debug, Clone, Deserialize)]
43#[serde(default)]
44pub struct Source {
45    pub enabled: bool,
46    /// Log root for this adapter; `None` means the adapter's own default.
47    pub path: Option<PathBuf>,
48}
49
50impl Default for Source {
51    fn default() -> Self {
52        Self {
53            enabled: true,
54            path: None,
55        }
56    }
57}
58
59/// Per-million-token rates for one model. `cache_write` is optional because not
60/// every provider bills it separately.
61#[derive(Debug, Clone, Copy, Deserialize)]
62pub struct ModelPrice {
63    #[serde(default)]
64    pub input: f64,
65    #[serde(default)]
66    pub output: f64,
67    #[serde(default)]
68    pub cache_read: f64,
69    #[serde(default)]
70    pub cache_write: Option<f64>,
71}
72
73/// Token counts to price. Absent counts are `None`, never `0`.
74#[derive(Debug, Clone, Copy, Default)]
75pub struct TokenCounts {
76    pub input: Option<u64>,
77    pub output: Option<u64>,
78    pub cache_read: Option<u64>,
79    pub cache_write: Option<u64>,
80}
81
82impl Config {
83    /// Read `config.toml` from `dir`. A missing file is not an error — it means
84    /// "all defaults".
85    pub fn load_from_dir(dir: &Path) -> Result<Self, ConfigError> {
86        Self::load_file(&dir.join("config.toml"))
87    }
88
89    /// Read a specific config file. A missing file yields the defaults.
90    pub fn load_file(path: &Path) -> Result<Self, ConfigError> {
91        let text = match std::fs::read_to_string(path) {
92            Ok(text) => text,
93            Err(err) if err.kind() == io::ErrorKind::NotFound => return Ok(Self::default()),
94            Err(err) => return Err(ConfigError::Read(path.to_path_buf(), err)),
95        };
96        toml::from_str(&text).map_err(|err| ConfigError::Parse(path.to_path_buf(), err))
97    }
98
99    /// Configured source settings, or the defaults for an unconfigured adapter.
100    pub fn source(&self, adapter: &str) -> Source {
101        self.sources.get(adapter).cloned().unwrap_or_default()
102    }
103
104    /// Configured price for a model, if any.
105    pub fn price(&self, provider: &str, model: &str) -> Option<ModelPrice> {
106        self.pricing.get(provider)?.get(model).copied()
107    }
108
109    /// Estimate cost in whole currency units for the given token counts.
110    ///
111    /// Returns `None` when the model has no configured price — callers must
112    /// propagate the absence rather than substituting `0.0`.
113    pub fn estimate_cost(&self, provider: &str, model: &str, tokens: TokenCounts) -> Option<f64> {
114        self.pricing().estimate_cost(provider, model, tokens)
115    }
116
117    /// The price table on its own, detached from the rest of the config.
118    ///
119    /// Reports price at *read* time, so they carry this rather than a borrow of
120    /// the whole config: editing `config.toml` re-prices events that are already
121    /// in the store, without a re-ingest.
122    pub fn pricing(&self) -> Pricing {
123        Pricing {
124            table: self.pricing.clone(),
125        }
126    }
127}
128
129/// A price table, owned. Empty by default, and an empty table prices nothing.
130#[derive(Debug, Clone, Default)]
131pub struct Pricing {
132    table: BTreeMap<String, BTreeMap<String, ModelPrice>>,
133}
134
135impl Pricing {
136    /// Configured price for a model, if any.
137    pub fn price(&self, provider: &str, model: &str) -> Option<ModelPrice> {
138        self.table.get(provider)?.get(model).copied()
139    }
140
141    /// `None` when this model has no configured price — never `0.0`.
142    pub fn estimate_cost(&self, provider: &str, model: &str, tokens: TokenCounts) -> Option<f64> {
143        let price = self.price(provider, model)?;
144        let per_million = |count: Option<u64>, rate: f64| count.unwrap_or(0) as f64 * rate / 1e6;
145        // An unset cache_write rate falls back to the input rate, matching how
146        // providers that do not bill writes separately behave.
147        let cache_write_rate = price.cache_write.unwrap_or(price.input);
148        Some(
149            per_million(tokens.input, price.input)
150                + per_million(tokens.output, price.output)
151                + per_million(tokens.cache_read, price.cache_read)
152                + per_million(tokens.cache_write, cache_write_rate),
153        )
154    }
155}
156
157#[derive(Debug)]
158pub enum ConfigError {
159    Read(PathBuf, io::Error),
160    Parse(PathBuf, toml::de::Error),
161}
162
163impl fmt::Display for ConfigError {
164    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
165        match self {
166            ConfigError::Read(path, err) => write!(f, "reading {}: {err}", path.display()),
167            ConfigError::Parse(path, err) => write!(f, "parsing {}: {err}", path.display()),
168        }
169    }
170}
171
172impl std::error::Error for ConfigError {}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    const SAMPLE: &str = r#"
179[general]
180index_prompt_text = false
181
182[sources.claude-code]
183enabled = true
184path = "/logs/claude"
185
186[pricing.anthropic]
187"claude-sonnet-4-6" = { input = 3.0, output = 15.0, cache_read = 0.3 }
188"#;
189
190    fn write(dir: &Path, text: &str) {
191        std::fs::write(dir.join("config.toml"), text).unwrap();
192    }
193
194    #[test]
195    fn missing_file_yields_defaults() {
196        let dir = tempfile::tempdir().unwrap();
197        let cfg = Config::load_from_dir(dir.path()).unwrap();
198        assert!(cfg.general.index_prompt_text);
199        assert!(cfg.general.data_dir.is_none());
200        assert!(cfg.pricing.is_empty());
201        // An unconfigured source is enabled with the adapter's own default path.
202        let source = cfg.source("claude-code");
203        assert!(source.enabled);
204        assert!(source.path.is_none());
205    }
206
207    #[test]
208    fn parses_all_sections() {
209        let dir = tempfile::tempdir().unwrap();
210        write(dir.path(), SAMPLE);
211        let cfg = Config::load_from_dir(dir.path()).unwrap();
212        assert!(!cfg.general.index_prompt_text);
213        assert_eq!(
214            cfg.source("claude-code").path.as_deref(),
215            Some(Path::new("/logs/claude"))
216        );
217        let price = cfg.price("anthropic", "claude-sonnet-4-6").unwrap();
218        assert_eq!(price.input, 3.0);
219        assert_eq!(price.cache_write, None);
220    }
221
222    #[test]
223    fn estimates_cost_from_config() {
224        let dir = tempfile::tempdir().unwrap();
225        write(dir.path(), SAMPLE);
226        let cfg = Config::load_from_dir(dir.path()).unwrap();
227        let tokens = TokenCounts {
228            input: Some(1_000_000),
229            output: Some(1_000_000),
230            cache_read: Some(1_000_000),
231            cache_write: None,
232        };
233        let cost = cfg
234            .estimate_cost("anthropic", "claude-sonnet-4-6", tokens)
235            .unwrap();
236        assert!((cost - 18.3).abs() < 1e-9, "got {cost}");
237    }
238
239    #[test]
240    fn unpriced_model_returns_none_not_zero() {
241        let dir = tempfile::tempdir().unwrap();
242        write(dir.path(), SAMPLE);
243        let cfg = Config::load_from_dir(dir.path()).unwrap();
244        let tokens = TokenCounts {
245            input: Some(1_000),
246            ..TokenCounts::default()
247        };
248        assert!(cfg
249            .estimate_cost("anthropic", "some-unlisted-model", tokens)
250            .is_none());
251        assert!(cfg.estimate_cost("openai", "gpt-x", tokens).is_none());
252    }
253
254    #[test]
255    fn empty_config_prices_nothing() {
256        let cfg = Config::default();
257        assert!(cfg
258            .estimate_cost("anthropic", "claude-sonnet-4-6", TokenCounts::default())
259            .is_none());
260    }
261}