Skip to main content

tokenmiser_config/
lib.rs

1//! Configuration types shared across every crate, kept here to avoid circular
2//! dependencies.
3
4use serde::{Deserialize, Serialize};
5use std::collections::HashMap;
6
7pub mod pricing;
8pub use pricing::{ModelPricing, PricingTable};
9
10/// Drives request shape and base-URL selection.
11#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
12pub enum ProviderKind {
13    #[serde(rename = "openai")]
14    OpenAI,
15    #[serde(rename = "anthropic")]
16    Anthropic,
17    #[serde(rename = "ollama")]
18    Ollama,
19    #[serde(rename = "deepseek")]
20    DeepSeek,
21    #[serde(rename = "gemini")]
22    Gemini,
23}
24
25impl ProviderKind {
26    pub fn as_str(&self) -> &'static str {
27        match self {
28            Self::OpenAI => "openai",
29            Self::Anthropic => "anthropic",
30            Self::Ollama => "ollama",
31            Self::DeepSeek => "deepseek",
32            Self::Gemini => "gemini",
33        }
34    }
35}
36
37/// A single configured upstream provider instance.
38#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct ProviderConfig {
40    pub name: String,
41    pub kind: ProviderKind,
42    /// Base URL including /v1 prefix where applicable.
43    pub base_url: String,
44    /// Name of the env var holding the API key; read at runtime, never logged.
45    #[serde(default)]
46    pub api_key_env: Option<String>,
47    /// Default model used when the request specifies a generic name.
48    #[serde(default)]
49    pub default_model: Option<String>,
50}
51
52impl ProviderConfig {
53    pub fn openai() -> Self {
54        Self {
55            name: "openai".into(),
56            kind: ProviderKind::OpenAI,
57            base_url: "https://api.openai.com/v1".into(),
58            api_key_env: Some("OPENAI_API_KEY".into()),
59            default_model: Some("gpt-5".into()),
60        }
61    }
62
63    pub fn anthropic() -> Self {
64        Self {
65            name: "anthropic".into(),
66            kind: ProviderKind::Anthropic,
67            base_url: "https://api.anthropic.com/v1".into(),
68            api_key_env: Some("ANTHROPIC_API_KEY".into()),
69            default_model: Some("claude-sonnet-4-6".into()),
70        }
71    }
72
73    pub fn ollama_local() -> Self {
74        Self {
75            name: "ollama".into(),
76            kind: ProviderKind::Ollama,
77            base_url: "http://localhost:11434".into(),
78            api_key_env: None,
79            default_model: Some("llama3.2".into()),
80        }
81    }
82}
83
84/// Top-level TokenMiser daemon config.
85#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct TokenmiserConfig {
87    pub listen: ListenConfig,
88    pub providers: Vec<ProviderConfig>,
89    #[serde(default)]
90    pub routing: RoutingConfig,
91    #[serde(default)]
92    pub cache: CacheConfig,
93    #[serde(default)]
94    pub budget: BudgetConfig,
95    #[serde(default)]
96    pub security: SecurityConfig,
97}
98
99impl Default for TokenmiserConfig {
100    fn default() -> Self {
101        Self {
102            listen: ListenConfig::default(),
103            providers: vec![
104                ProviderConfig::openai(),
105                ProviderConfig::anthropic(),
106                ProviderConfig::ollama_local(),
107            ],
108            routing: RoutingConfig::default(),
109            cache: CacheConfig::default(),
110            budget: BudgetConfig::default(),
111            security: SecurityConfig::default(),
112        }
113    }
114}
115
116/// Browser-facing defenses for the local proxy.
117///
118/// Loopback is not a boundary against the browser: any page the operator
119/// visits can POST to 127.0.0.1 as a CORS simple request, which is sent
120/// without a preflight. CORS blocks reading the reply, but the request has
121/// already spent the budget and can poison the cache.
122#[derive(Debug, Clone, Default, Serialize, Deserialize)]
123pub struct SecurityConfig {
124    /// Origins permitted to drive the proxy from a browser, matched exactly
125    /// against the `Origin` header. Empty (the default) accepts no
126    /// cross-origin browser traffic; non-browser clients send no such header
127    /// and are unaffected either way.
128    ///
129    /// The literal `"*"` disables the check entirely, re-opening the CSRF hole
130    /// for every page the operator visits.
131    #[serde(default)]
132    pub allowed_origins: Vec<String>,
133}
134
135impl SecurityConfig {
136    /// True when `origin` may drive the proxy from a browser.
137    pub fn origin_allowed(&self, origin: &str) -> bool {
138        self.allowed_origins.iter().any(|o| o == "*" || o == origin)
139    }
140
141    /// True when the operator has opened the proxy to every origin, which is
142    /// the only case where browser-origin signals are ignored.
143    pub fn allows_any_origin(&self) -> bool {
144        self.allowed_origins.iter().any(|o| o == "*")
145    }
146}
147
148/// Spend-budget thresholds. A crossed limit is surfaced via `/stats`, the
149/// `x-tokenmiser-budget` header and a log line, but blocks requests only under
150/// `enforce`, and then only paid routes.
151#[derive(Debug, Clone, Default, Serialize, Deserialize)]
152pub struct BudgetConfig {
153    /// Max USD spend per UTC day.
154    #[serde(default)]
155    pub daily_usd: Option<f64>,
156    /// Max USD spend over the daemon's lifetime.
157    #[serde(default)]
158    pub total_usd: Option<f64>,
159    /// Reject paid-provider requests with HTTP 402 once a limit is exceeded.
160    #[serde(default)]
161    pub enforce: bool,
162}
163
164impl BudgetConfig {
165    pub fn is_active(&self) -> bool {
166        self.daily_usd.is_some() || self.total_usd.is_some()
167    }
168}
169
170#[derive(Debug, Clone, Serialize, Deserialize)]
171pub struct ListenConfig {
172    /// Pingora proxy ingress: the LLM-traffic surface.
173    pub proxy_addr: String,
174    /// Admin ingress: `/stats`, `/healthz`, dashboard.
175    pub admin_addr: String,
176}
177
178impl Default for ListenConfig {
179    fn default() -> Self {
180        // The proxy is unauthenticated and can spend the operator's API
181        // budget, so LAN exposure must be an explicit opt-in.
182        Self {
183            proxy_addr: "127.0.0.1:8443".into(),
184            admin_addr: "127.0.0.1:9443".into(),
185        }
186    }
187}
188
189#[derive(Debug, Clone, Default, Serialize, Deserialize)]
190pub struct RoutingConfig {
191    /// Static model aliases, e.g. `gpt-5` always routing to `openai`.
192    #[serde(default)]
193    pub aliases: HashMap<String, ModelTarget>,
194    /// Default provider when the requested model is unrecognized.
195    #[serde(default)]
196    pub default_provider: Option<String>,
197}
198
199#[derive(Debug, Clone, Serialize, Deserialize)]
200pub struct ModelTarget {
201    pub provider: String,
202    pub model: String,
203}
204
205#[derive(Debug, Clone, Serialize, Deserialize)]
206pub struct CacheConfig {
207    #[serde(default = "default_true")]
208    pub l1_enabled: bool,
209    #[serde(default = "default_true")]
210    pub l2_enabled: bool,
211    #[serde(default = "default_semantic_threshold")]
212    pub semantic_threshold: f32,
213    /// Skip L2 candidates whose prompt carries a different set of number
214    /// literals than the query. Embeddings are near-blind to digits, so
215    /// without this a cached "Add 4 and 9" can answer "Multiply 3 by 11".
216    /// Disable only for workloads that want digit-insensitive matching.
217    #[serde(default = "default_true")]
218    pub numeric_guard: bool,
219    /// `tenant` | `user` | `session` | `global`
220    #[serde(default = "default_scope")]
221    pub scope: String,
222}
223
224impl Default for CacheConfig {
225    fn default() -> Self {
226        Self {
227            l1_enabled: true,
228            l2_enabled: true,
229            semantic_threshold: default_semantic_threshold(),
230            numeric_guard: true,
231            scope: "tenant".into(),
232        }
233    }
234}
235
236fn default_true() -> bool {
237    true
238}
239fn default_semantic_threshold() -> f32 {
240    // Tuned against `tokenmiser-cache::threshold_bench`. Biased toward
241    // precision: a false-positive cache hit returns a wrong answer silently.
242    0.87
243}
244fn default_scope() -> String {
245    "tenant".into()
246}
247
248#[cfg(test)]
249mod tests {
250    use super::*;
251
252    #[test]
253    fn budget_defaults_to_inactive_warn_only() {
254        let cfg = TokenmiserConfig::default();
255        assert!(!cfg.budget.is_active());
256        assert!(!cfg.budget.enforce);
257        assert!(cfg.budget.daily_usd.is_none());
258        assert!(cfg.budget.total_usd.is_none());
259    }
260
261    #[test]
262    fn budget_parses_from_yaml() {
263        let yaml = r#"
264listen:
265  proxy_addr: "127.0.0.1:1"
266  admin_addr: "127.0.0.1:2"
267providers: []
268budget:
269  daily_usd: 5.0
270  enforce: true
271"#;
272        let cfg: TokenmiserConfig = serde_yaml::from_str(yaml).unwrap();
273        assert!(cfg.budget.is_active());
274        assert_eq!(cfg.budget.daily_usd, Some(5.0));
275        assert_eq!(cfg.budget.total_usd, None);
276        assert!(cfg.budget.enforce);
277    }
278
279    #[test]
280    fn numeric_guard_defaults_on_and_is_config_exposed() {
281        assert!(CacheConfig::default().numeric_guard);
282
283        let yaml = r#"
284listen:
285  proxy_addr: "127.0.0.1:1"
286  admin_addr: "127.0.0.1:2"
287providers: []
288cache:
289  numeric_guard: false
290  semantic_threshold: 0.91
291"#;
292        let cfg: TokenmiserConfig = serde_yaml::from_str(yaml).unwrap();
293        assert!(!cfg.cache.numeric_guard);
294        assert_eq!(cfg.cache.semantic_threshold, 0.91);
295
296        // Omitting the field keeps the safe default.
297        let yaml = r#"
298listen:
299  proxy_addr: "127.0.0.1:1"
300  admin_addr: "127.0.0.1:2"
301providers: []
302cache:
303  semantic_threshold: 0.87
304"#;
305        let cfg: TokenmiserConfig = serde_yaml::from_str(yaml).unwrap();
306        assert!(cfg.cache.numeric_guard);
307    }
308
309    #[test]
310    fn security_defaults_to_no_allowed_origins() {
311        let cfg = TokenmiserConfig::default();
312        assert!(cfg.security.allowed_origins.is_empty());
313        assert!(!cfg.security.allows_any_origin());
314        assert!(!cfg.security.origin_allowed("https://evil.example"));
315        assert!(!cfg.security.origin_allowed("http://localhost:3000"));
316    }
317
318    #[test]
319    fn security_allowed_origins_parse_and_match_exactly() {
320        let yaml = r#"
321listen:
322  proxy_addr: "127.0.0.1:1"
323  admin_addr: "127.0.0.1:2"
324providers: []
325security:
326  allowed_origins:
327    - "http://localhost:3000"
328    - "https://app.example.com"
329"#;
330        let cfg: TokenmiserConfig = serde_yaml::from_str(yaml).unwrap();
331        assert!(cfg.security.origin_allowed("http://localhost:3000"));
332        assert!(cfg.security.origin_allowed("https://app.example.com"));
333        assert!(!cfg.security.origin_allowed("http://localhost:3001"));
334        assert!(!cfg.security.origin_allowed("https://localhost:3000"));
335        assert!(!cfg
336            .security
337            .origin_allowed("https://app.example.com.evil.test"));
338        assert!(!cfg.security.allows_any_origin());
339    }
340
341    #[test]
342    fn security_wildcard_opens_every_origin() {
343        let yaml = r#"
344listen:
345  proxy_addr: "127.0.0.1:1"
346  admin_addr: "127.0.0.1:2"
347providers: []
348security:
349  allowed_origins: ["*"]
350"#;
351        let cfg: TokenmiserConfig = serde_yaml::from_str(yaml).unwrap();
352        assert!(cfg.security.allows_any_origin());
353        assert!(cfg.security.origin_allowed("https://evil.example"));
354    }
355
356    #[test]
357    fn config_without_budget_section_still_parses() {
358        let yaml = r#"
359listen:
360  proxy_addr: "127.0.0.1:1"
361  admin_addr: "127.0.0.1:2"
362providers: []
363"#;
364        let cfg: TokenmiserConfig = serde_yaml::from_str(yaml).unwrap();
365        assert!(!cfg.budget.is_active());
366    }
367}