Skip to main content

codei_config/
model.rs

1use std::collections::HashMap;
2use std::path::PathBuf;
3
4use serde::{Deserialize, Serialize};
5
6/// Full configuration after merging all sources.
7#[derive(Debug, Clone, Serialize, Deserialize)]
8pub struct Config {
9    #[serde(default)]
10    pub defaults: DefaultsConfig,
11
12    #[serde(default)]
13    pub providers: HashMap<String, ProviderConfig>,
14
15    #[serde(default)]
16    pub agent: AgentConfig,
17
18    #[serde(default)]
19    pub tools: ToolsConfig,
20
21    #[serde(default)]
22    pub ui: UiConfig,
23
24    #[serde(default)]
25    pub session: SessionConfig,
26}
27
28/// Resolved configuration with metadata about where values came from.
29#[derive(Debug, Clone)]
30pub struct ResolvedConfig {
31    pub config: Config,
32    pub cwd: PathBuf,
33    pub project_root: Option<PathBuf>,
34    pub user_config_path: PathBuf,
35    pub project_config_path: Option<PathBuf>,
36}
37
38#[derive(Debug, Clone, Serialize, Deserialize)]
39pub struct DefaultsConfig {
40    pub model: String,
41    pub provider: String,
42    pub temperature: f32,
43    pub max_tokens: u32,
44    pub language: String,
45}
46
47#[derive(Debug, Clone, Serialize, Deserialize)]
48pub struct ProviderConfig {
49    /// Direct API key; takes precedence over `api_key_env` when set.
50    #[serde(default)]
51    pub api_key: Option<String>,
52    #[serde(default)]
53    pub api_key_env: Option<String>,
54    #[serde(default)]
55    pub base_url: Option<String>,
56    #[serde(default)]
57    pub api_style: Option<String>,
58    /// `tools` (OpenAI tools API) or `functions` (legacy function calling). Default: `functions`.
59    #[serde(default)]
60    pub tool_format: Option<String>,
61}
62
63impl ProviderConfig {
64    /// Resolve API key: `api_key` first, then the environment variable named by `api_key_env`.
65    pub fn resolve_api_key(&self) -> Result<String, crate::ConfigError> {
66        if let Some(key) = self
67            .api_key
68            .as_deref()
69            .map(str::trim)
70            .filter(|k| !k.is_empty())
71        {
72            return Ok(key.to_string());
73        }
74        let env_name = self
75            .api_key_env
76            .as_deref()
77            .filter(|name| !name.is_empty())
78            .unwrap_or("OPENAI_API_KEY");
79        std::env::var(env_name).map_err(|_| crate::ConfigError::MissingApiKey {
80            env: env_name.to_string(),
81        })
82    }
83}
84
85#[derive(Debug, Clone, Serialize, Deserialize)]
86pub struct AgentConfig {
87    pub max_turns: u32,
88    pub max_tool_rounds_per_turn: u32,
89    pub context_window_tokens: u32,
90    /// Fraction of `context_window_tokens` that triggers automatic session compaction (0.0–1.0).
91    pub compaction_threshold: f32,
92    /// Number of recent session messages to keep after compaction.
93    pub compaction_keep_messages: u32,
94    /// Max tokens for the LLM-generated compaction summary.
95    pub compaction_summary_max_tokens: u32,
96}
97
98#[derive(Debug, Clone, Serialize, Deserialize)]
99pub struct ToolsConfig {
100    pub shell: ShellToolConfig,
101    pub write: EnabledToolConfig,
102    #[serde(default)]
103    pub web_search: WebSearchToolConfig,
104    #[serde(default)]
105    pub web_fetch: WebFetchToolConfig,
106    #[serde(default)]
107    pub grep: GrepToolConfig,
108}
109
110#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
111#[serde(rename_all = "snake_case")]
112pub enum WebSearchProvider {
113    #[default]
114    Duckduckgo,
115    Searxng,
116}
117
118#[derive(Debug, Clone, Serialize, Deserialize)]
119pub struct WebSearchToolConfig {
120    pub enabled: bool,
121    pub provider: WebSearchProvider,
122    pub timeout_secs: u64,
123    pub max_results: usize,
124    #[serde(default)]
125    pub searxng_url: Option<String>,
126    #[serde(default = "default_ssrf_protection")]
127    pub ssrf_protection: bool,
128}
129
130#[derive(Debug, Clone, Serialize, Deserialize)]
131pub struct WebFetchToolConfig {
132    pub enabled: bool,
133    pub timeout_secs: u64,
134    pub max_bytes: usize,
135    /// When true, block localhost, private IPs, and link-local addresses.
136    #[serde(default = "default_ssrf_protection")]
137    pub ssrf_protection: bool,
138}
139
140#[derive(Debug, Clone, Serialize, Deserialize)]
141pub struct GrepToolConfig {
142    pub max_matches: usize,
143    pub max_files: usize,
144}
145
146#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
147#[serde(rename_all = "lowercase")]
148pub enum ShellSandboxMode {
149    #[default]
150    Off,
151    Restricted,
152}
153
154#[derive(Debug, Clone, Serialize, Deserialize)]
155pub struct ShellToolConfig {
156    pub enabled: bool,
157    pub timeout_secs: u64,
158    #[serde(default)]
159    pub sandbox: ShellSandboxMode,
160    #[serde(default)]
161    pub allowlist: Vec<String>,
162}
163
164#[derive(Debug, Clone, Serialize, Deserialize)]
165pub struct EnabledToolConfig {
166    pub enabled: bool,
167}
168
169#[derive(Debug, Clone, Serialize, Deserialize)]
170pub struct UiConfig {
171    pub theme: UiTheme,
172    pub show_tool_output: bool,
173    pub confirm_destructive: bool,
174}
175
176#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
177#[serde(rename_all = "lowercase")]
178pub enum UiTheme {
179    #[default]
180    Auto,
181    Dark,
182    Light,
183}
184
185#[derive(Debug, Clone, Serialize, Deserialize)]
186pub struct SessionConfig {
187    pub storage: SessionStorage,
188    pub dir: String,
189}
190
191#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq, Default)]
192#[serde(rename_all = "lowercase")]
193pub enum SessionStorage {
194    #[default]
195    Sqlite,
196    Json,
197}
198
199impl Default for Config {
200    fn default() -> Self {
201        Self {
202            defaults: DefaultsConfig::default(),
203            providers: default_providers(),
204            agent: AgentConfig::default(),
205            tools: ToolsConfig::default(),
206            ui: UiConfig::default(),
207            session: SessionConfig::default(),
208        }
209    }
210}
211
212impl Default for DefaultsConfig {
213    fn default() -> Self {
214        Self {
215            model: "gpt-4o".to_string(),
216            provider: "openai".to_string(),
217            temperature: 0.2,
218            max_tokens: 8192,
219            language: "zh-CN".to_string(),
220        }
221    }
222}
223
224impl Default for AgentConfig {
225    fn default() -> Self {
226        Self {
227            max_turns: 50,
228            max_tool_rounds_per_turn: 25,
229            context_window_tokens: 128_000,
230            compaction_threshold: 0.85,
231            compaction_keep_messages: 12,
232            compaction_summary_max_tokens: 2048,
233        }
234    }
235}
236
237impl Default for ToolsConfig {
238    fn default() -> Self {
239        Self {
240            shell: ShellToolConfig {
241                enabled: true,
242                timeout_secs: 120,
243                sandbox: ShellSandboxMode::Off,
244                allowlist: Vec::new(),
245            },
246            write: EnabledToolConfig { enabled: true },
247            web_search: WebSearchToolConfig::default(),
248            web_fetch: WebFetchToolConfig::default(),
249            grep: GrepToolConfig::default(),
250        }
251    }
252}
253
254impl Default for WebSearchToolConfig {
255    fn default() -> Self {
256        Self {
257            enabled: false,
258            provider: WebSearchProvider::Duckduckgo,
259            timeout_secs: 30,
260            max_results: 10,
261            searxng_url: None,
262            ssrf_protection: true,
263        }
264    }
265}
266
267impl Default for WebFetchToolConfig {
268    fn default() -> Self {
269        Self {
270            enabled: true,
271            timeout_secs: 30,
272            max_bytes: 512 * 1024,
273            ssrf_protection: true,
274        }
275    }
276}
277
278fn default_ssrf_protection() -> bool {
279    true
280}
281
282impl Default for GrepToolConfig {
283    fn default() -> Self {
284        Self {
285            max_matches: 200,
286            max_files: 5_000,
287        }
288    }
289}
290
291impl Default for UiConfig {
292    fn default() -> Self {
293        Self {
294            theme: UiTheme::Auto,
295            show_tool_output: true,
296            confirm_destructive: true,
297        }
298    }
299}
300
301impl Default for SessionConfig {
302    fn default() -> Self {
303        Self {
304            storage: SessionStorage::Sqlite,
305            dir: "~/.local/share/codei/sessions".to_string(),
306        }
307    }
308}
309
310fn default_providers() -> HashMap<String, ProviderConfig> {
311    HashMap::from([
312        (
313            "openai".to_string(),
314            ProviderConfig {
315                api_key: None,
316                api_key_env: Some("OPENAI_API_KEY".to_string()),
317                base_url: Some("https://api.openai.com/v1".to_string()),
318                api_style: Some("openai".to_string()),
319                tool_format: Some("tools".to_string()),
320            },
321        ),
322        (
323            "anthropic".to_string(),
324            ProviderConfig {
325                api_key: None,
326                api_key_env: Some("ANTHROPIC_API_KEY".to_string()),
327                base_url: None,
328                api_style: Some("anthropic".to_string()),
329                tool_format: None,
330            },
331        ),
332        (
333            "custom".to_string(),
334            ProviderConfig {
335                api_key: None,
336                api_key_env: Some("CUSTOM_API_KEY".to_string()),
337                base_url: Some("http://localhost:8080/v1".to_string()),
338                api_style: Some("openai".to_string()),
339                tool_format: Some("tools".to_string()),
340            },
341        ),
342    ])
343}
344
345impl ResolvedConfig {
346    pub fn validate(&self) -> Result<(), crate::ConfigError> {
347        let lang = &self.config.defaults.language;
348        if lang != "zh-CN" && lang != "en-US" {
349            return Err(crate::ConfigError::InvalidLanguage {
350                language: lang.clone(),
351            });
352        }
353        let threshold = self.config.agent.compaction_threshold;
354        if !(0.0..=1.0).contains(&threshold) {
355            return Err(crate::ConfigError::InvalidCompactionThreshold { threshold });
356        }
357        Ok(())
358    }
359}
360
361#[cfg(test)]
362mod provider_tests {
363    use super::ProviderConfig;
364
365    #[test]
366    fn resolve_api_key_prefers_direct_config() {
367        let cfg = ProviderConfig {
368            api_key: Some("sk-direct".into()),
369            api_key_env: Some("OPENAI_API_KEY".into()),
370            base_url: None,
371            api_style: None,
372            tool_format: None,
373        };
374        assert_eq!(cfg.resolve_api_key().unwrap(), "sk-direct");
375    }
376
377    #[test]
378    fn blank_api_key_falls_back_to_env() {
379        let cfg = ProviderConfig {
380            api_key: Some("   ".into()),
381            api_key_env: Some("NONEXISTENT_CODEI_API_KEY".into()),
382            base_url: None,
383            api_style: None,
384            tool_format: None,
385        };
386        let err = cfg.resolve_api_key().unwrap_err();
387        assert!(matches!(
388            err,
389            crate::ConfigError::MissingApiKey {
390                env: ref name
391            } if name == "NONEXISTENT_CODEI_API_KEY"
392        ));
393    }
394}