Skip to main content

chub_core/
config.rs

1use serde::{Deserialize, Serialize};
2use std::path::PathBuf;
3
4const DEFAULT_CDN_URL: &str = "https://cdn.aichub.org/v1";
5/// Empty by default — online telemetry forwarding is opt-in.
6/// Set `telemetry_url` in config or CHUB_TELEMETRY_URL env var to enable.
7const DEFAULT_TELEMETRY_URL: &str = "";
8
9/// Maximum size for YAML config files (1 MB) to prevent denial-of-service via
10/// anchor bombs or deeply nested structures.
11const MAX_CONFIG_FILE_SIZE: u64 = 1024 * 1024;
12
13#[derive(Debug, Clone, Serialize, Deserialize)]
14pub struct SourceConfig {
15    pub name: String,
16    #[serde(skip_serializing_if = "Option::is_none")]
17    pub url: Option<String>,
18    #[serde(skip_serializing_if = "Option::is_none")]
19    pub path: Option<String>,
20}
21
22#[derive(Debug, Clone, Default, Serialize, Deserialize)]
23pub struct TrackingConfig {
24    /// Custom cost rates keyed by model name substring (e.g. "opus", "gpt-4o").
25    /// Each rate is [input_per_m, output_per_m] or [input, output, cache_read, cache_write].
26    #[serde(default, skip_serializing_if = "Vec::is_empty")]
27    pub cost_rates: Vec<CustomCostRate>,
28    /// Monthly budget alert threshold in USD (0 = disabled).
29    #[serde(default)]
30    pub budget_alert_usd: f64,
31}
32
33#[derive(Debug, Clone, Serialize, Deserialize)]
34pub struct CustomCostRate {
35    /// Model name substring to match (e.g. "opus", "gpt-4o", "my-fine-tune").
36    pub model: String,
37    /// Cost per million input tokens.
38    pub input_per_m: f64,
39    /// Cost per million output tokens.
40    pub output_per_m: f64,
41    /// Cost per million cache-read tokens (defaults to input * 0.1).
42    #[serde(default)]
43    pub cache_read_per_m: Option<f64>,
44    /// Cost per million cache-write tokens (defaults to input * 1.25).
45    #[serde(default)]
46    pub cache_write_per_m: Option<f64>,
47}
48
49#[derive(Debug, Clone)]
50pub struct Config {
51    pub sources: Vec<SourceConfig>,
52    pub output_dir: String,
53    pub refresh_interval: u64,
54    pub output_format: String,
55    pub source: String,
56    pub telemetry: bool,
57    pub feedback: bool,
58    pub telemetry_url: String,
59    pub annotation_token: Option<String>,
60    pub tracking: TrackingConfig,
61}
62
63impl Default for Config {
64    fn default() -> Self {
65        Self {
66            sources: vec![SourceConfig {
67                name: "default".to_string(),
68                url: Some(DEFAULT_CDN_URL.to_string()),
69                path: None,
70            }],
71            output_dir: ".context".to_string(),
72            refresh_interval: 21600,
73            output_format: "human".to_string(),
74            source: "official,maintainer,community".to_string(),
75            telemetry: true,
76            feedback: true,
77            telemetry_url: DEFAULT_TELEMETRY_URL.to_string(),
78            annotation_token: None,
79            tracking: TrackingConfig::default(),
80        }
81    }
82}
83
84/// Raw YAML config file structure.
85#[derive(Debug, Deserialize, Default)]
86struct FileConfig {
87    #[serde(default)]
88    sources: Option<Vec<SourceConfig>>,
89    #[serde(default)]
90    cdn_url: Option<String>,
91    #[serde(default)]
92    output_dir: Option<String>,
93    #[serde(default)]
94    refresh_interval: Option<u64>,
95    #[serde(default)]
96    output_format: Option<String>,
97    #[serde(default)]
98    source: Option<String>,
99    #[serde(default)]
100    telemetry: Option<bool>,
101    #[serde(default)]
102    feedback: Option<bool>,
103    #[serde(default)]
104    telemetry_url: Option<String>,
105    #[serde(default)]
106    annotation_token: Option<String>,
107    #[serde(default)]
108    tracking: Option<TrackingConfig>,
109}
110
111/// Get the chub data directory (~/.chub or CHUB_DIR env var).
112pub fn chub_dir() -> PathBuf {
113    if let Ok(dir) = std::env::var("CHUB_DIR") {
114        PathBuf::from(dir)
115    } else {
116        dirs::home_dir()
117            .unwrap_or_else(|| PathBuf::from("."))
118            .join(".chub")
119    }
120}
121
122/// Load config with three-tier merge:
123/// Tier 1: ~/.chub/config.yaml (personal defaults)
124/// Tier 2: .chub/config.yaml (project, overrides personal)
125/// Tier 3: Active profile (additive — doesn't override config fields)
126pub fn load_config() -> Config {
127    let defaults = Config::default();
128    let config_path = chub_dir().join("config.yaml");
129
130    // Tier 1: personal config (with size limit)
131    let file_config: FileConfig = read_config_file(&config_path).unwrap_or_default();
132
133    // Tier 2: project config (override personal, with size limit)
134    let project_config: FileConfig = find_project_chub_dir()
135        .map(|d| d.join("config.yaml"))
136        .and_then(|p| read_config_file(&p))
137        .unwrap_or_default();
138
139    // Merge: project overrides personal, personal overrides defaults
140    let sources = project_config
141        .sources
142        .or(file_config.sources)
143        .unwrap_or_else(|| {
144            let raw_url = std::env::var("CHUB_BUNDLE_URL")
145                .ok()
146                .or(project_config.cdn_url)
147                .or(file_config.cdn_url)
148                .unwrap_or_else(|| DEFAULT_CDN_URL.to_string());
149            // Validate the URL scheme (HTTPS required, HTTP only for localhost)
150            let url = match crate::util::validate_url(&raw_url, "CHUB_BUNDLE_URL") {
151                Ok(u) => u,
152                Err(e) => {
153                    eprintln!("Warning: {}", e);
154                    DEFAULT_CDN_URL.to_string()
155                }
156            };
157            vec![SourceConfig {
158                name: "default".to_string(),
159                url: Some(url),
160                path: None,
161            }]
162        });
163
164    Config {
165        sources,
166        output_dir: project_config
167            .output_dir
168            .or(file_config.output_dir)
169            .unwrap_or(defaults.output_dir),
170        refresh_interval: project_config
171            .refresh_interval
172            .or(file_config.refresh_interval)
173            .unwrap_or(defaults.refresh_interval),
174        output_format: project_config
175            .output_format
176            .or(file_config.output_format)
177            .unwrap_or(defaults.output_format),
178        source: project_config
179            .source
180            .or(file_config.source)
181            .unwrap_or(defaults.source),
182        telemetry: project_config
183            .telemetry
184            .or(file_config.telemetry)
185            .unwrap_or(defaults.telemetry),
186        feedback: project_config
187            .feedback
188            .or(file_config.feedback)
189            .unwrap_or(defaults.feedback),
190        telemetry_url: project_config
191            .telemetry_url
192            .or(file_config.telemetry_url)
193            .unwrap_or(defaults.telemetry_url),
194        // annotation_token intentionally comes from personal config only —
195        // never from project config to avoid accidental token commits.
196        annotation_token: file_config.annotation_token,
197        tracking: project_config
198            .tracking
199            .or(file_config.tracking)
200            .unwrap_or_default(),
201    }
202}
203
204/// Get the annotation server auth token.
205/// Priority: CHUB_ANNOTATION_TOKEN env var > ~/.chub/config.yaml annotation_token.
206/// Token is intentionally NOT read from .chub/config.yaml to prevent accidental commits.
207pub fn get_annotation_token() -> Option<String> {
208    std::env::var("CHUB_ANNOTATION_TOKEN")
209        .ok()
210        .or_else(|| load_config().annotation_token)
211}
212
213/// Read and parse a YAML config file with a size limit.
214fn read_config_file(path: &std::path::Path) -> Option<FileConfig> {
215    let meta = std::fs::metadata(path).ok()?;
216    if meta.len() > MAX_CONFIG_FILE_SIZE {
217        eprintln!(
218            "Warning: config file {} exceeds size limit ({} bytes, max {}). Skipping.",
219            path.display(),
220            meta.len(),
221            MAX_CONFIG_FILE_SIZE
222        );
223        return None;
224    }
225    let raw = std::fs::read_to_string(path).ok()?;
226    serde_yaml::from_str(&raw).ok()
227}
228
229/// Search upward from CWD for a `.chub/` directory (project-level).
230fn find_project_chub_dir() -> Option<PathBuf> {
231    let start = std::env::current_dir().ok()?;
232    let mut current = start.as_path();
233    loop {
234        let candidate = current.join(".chub");
235        if candidate.is_dir() {
236            // Don't use ~/.chub as a project dir
237            let home_chub = chub_dir();
238            if candidate != home_chub {
239                return Some(candidate);
240            }
241        }
242        current = current.parent()?;
243    }
244}