robit_ai/
config.rs

1//! Configuration loading for config.toml.
2//!
3//! Loads a single unified config file from:
4//!   1. `cwd/.robit/config.toml` (project-local, highest priority)
5//!   2. `~/.robit/config.toml`   (global fallback)
6//!
7//! Configuration format uses a providers + models structure:
8//! ```toml
9//! default_model = "deepseek/deepseek-chat"
10//!
11//! [providers.deepseek]
12//! name = "DeepSeek"
13//! base_url = "https://api.deepseek.com/v1"
14//! api_key = "${DEEPSEEK_API_KEY}"
15//!
16//! [[providers.deepseek.models]]
17//! id = "deepseek-chat"
18//! context_window = 65536
19//! ```
20//!
21//! Environment variable substitution is supported in `api_key` fields via `${ENV_VAR}` syntax.
22
23use serde::Deserialize;
24use std::collections::HashMap;
25use std::path::PathBuf;
26
27use crate::error::LlmError;
28
29// ============================================================================
30// config.toml structures
31// ============================================================================
32
33/// Top-level config.toml configuration.
34#[derive(Debug, Deserialize)]
35pub struct RobitConfig {
36    /// Default model in "provider/model" format (e.g. "deepseek/deepseek-chat").
37    pub default_model: Option<String>,
38    /// Provider definitions keyed by provider name.
39    pub providers: HashMap<String, ProviderConfig>,
40    /// Application settings.
41    pub app: Option<AppConfig>,
42}
43
44/// A single LLM provider (one API endpoint with multiple models).
45#[derive(Debug, Deserialize)]
46pub struct ProviderConfig {
47    /// Display name for the provider (optional).
48    pub name: Option<String>,
49    /// API base URL (must be OpenAI-compatible).
50    pub base_url: String,
51    /// API key (supports `${ENV_VAR}` substitution).
52    pub api_key: String,
53    /// Available models under this provider.
54    pub models: Vec<ModelConfig>,
55}
56
57/// A single model definition within a provider.
58#[derive(Debug, Deserialize)]
59pub struct ModelConfig {
60    /// Model ID used in API calls (e.g. "deepseek-chat").
61    pub id: String,
62    /// Display name (optional).
63    pub name: Option<String>,
64    /// Context window size in tokens (optional).
65    pub context_window: Option<u64>,
66    /// Maximum output tokens (optional).
67    pub max_output_tokens: Option<u64>,
68    /// Sampling temperature (optional, runtime parameter).
69    pub temperature: Option<f32>,
70    /// Maximum completion tokens (optional, runtime parameter).
71    pub max_tokens: Option<u32>,
72    /// Whether this model supports image inputs (optional, default false).
73    pub supports_images: Option<bool>,
74    /// Whether this model supports tool calling (optional, default false).
75    pub supports_tools: Option<bool>,
76}
77
78// ============================================================================
79// Application config (unchanged from previous version)
80// ============================================================================
81
82#[derive(Debug, Deserialize, Default)]
83pub struct AppConfig {
84    pub log_level: Option<String>,
85    pub max_steps: Option<usize>,
86    pub enabled_tools: Option<Vec<String>>,
87    pub enabled_skills: Option<Vec<String>>,
88    pub context: Option<ContextConfig>,
89    pub retry: Option<RetryConfig>,
90    pub auto_approve: Option<bool>,
91}
92
93#[derive(Debug, Deserialize)]
94pub struct ContextConfig {
95    pub max_output_lines: Option<usize>,
96    pub max_output_bytes: Option<usize>,
97    pub reserve_ratio: Option<f32>,
98    /// Token threshold for triggering compression (default 5000).
99    /// Only compress when removed messages exceed this token count.
100    pub compression_token_threshold: Option<usize>,
101    /// Enable/disable context compression (default true).
102    pub compression_enabled: Option<bool>,
103}
104
105#[derive(Debug, Deserialize)]
106pub struct RetryConfig {
107    pub max_retries: Option<u32>,
108    pub initial_backoff_ms: Option<u64>,
109    pub max_backoff_ms: Option<u64>,
110}
111
112// ============================================================================
113// Resolved model reference
114// ============================================================================
115
116/// A fully resolved model ready for client construction.
117///
118/// Merges provider-level settings (base_url, api_key) with model-level
119/// settings (context_window, temperature, etc).
120#[derive(Debug, Clone)]
121pub struct ResolvedModel {
122    pub profile_name: String,
123    pub model_id: String,
124    pub base_url: String,
125    pub api_key: String,
126    pub max_tokens: Option<u32>,
127    pub temperature: Option<f32>,
128    pub context_window: Option<u64>,
129}
130
131// ============================================================================
132// Loader
133// ============================================================================
134
135/// Returns the ~/.robit/ directory path.
136fn robit_home() -> Result<PathBuf, LlmError> {
137    let home = dirs::home_dir().ok_or_else(|| {
138        LlmError::ConfigError("Cannot determine home directory".to_string())
139    })?;
140    Ok(home.join(".robit"))
141}
142
143/// Replace `${ENV_VAR}` patterns with actual environment variable values.
144fn resolve_env_var(value: &str) -> String {
145    if let Some(var_name) = value.strip_prefix("${").and_then(|s| s.strip_suffix('}')) {
146        std::env::var(var_name).unwrap_or_else(|_| value.to_string())
147    } else {
148        value.to_string()
149    }
150}
151
152/// Load and parse the config.toml config file.
153///
154/// Automatically loads `~/.robit/.env` before resolving `${ENV_VAR}` patterns.
155///
156/// Search order:
157///   1. `workdir/.robit/config.toml` (project-local, if workdir provided)
158///   2. `cwd/.robit/config.toml` (project-local, if workdir not provided)
159///   3. `~/.robit/config.toml`   (global fallback)
160pub fn load_config(workdir: Option<&std::path::Path>) -> Result<RobitConfig, LlmError> {
161    // Load .env first so ${ENV_VAR} substitutions work
162    load_env();
163
164    let path = find_config_path(workdir)?;
165
166    let content = std::fs::read_to_string(&path).map_err(|e| {
167        LlmError::ConfigError(format!("Failed to read {}: {}", path.display(), e))
168    })?;
169
170    let mut config: RobitConfig = toml::from_str(&content).map_err(|e| {
171        LlmError::ConfigError(format!("Failed to parse config.toml: {}", e))
172    })?;
173
174    // Resolve environment variables in api_key fields
175    for provider in config.providers.values_mut() {
176        provider.api_key = resolve_env_var(&provider.api_key);
177    }
178
179    Ok(config)
180}
181
182/// Load .env from ~/.robit/.env if it exists.
183pub fn load_env() {
184    if let Ok(robit_dir) = robit_home() {
185        let env_path = robit_dir.join(".env");
186        if env_path.exists() {
187            let _ = dotenvy::from_path(&env_path);
188        }
189    }
190}
191
192/// Find the config file path following the search order.
193fn find_config_path(workdir: Option<&std::path::Path>) -> Result<PathBuf, LlmError> {
194    // 1. Project-local: workdir/.robit/config.toml (if workdir provided)
195    if let Some(workdir) = workdir {
196        let local_path = workdir.join(".robit").join("config.toml");
197        if local_path.exists() {
198            return Ok(local_path);
199        }
200    }
201
202    // 2. Project-local: cwd/.robit/config.toml (if workdir not provided or no config there)
203    if let Ok(cwd) = std::env::current_dir() {
204        let local_path = cwd.join(".robit").join("config.toml");
205        if local_path.exists() {
206            return Ok(local_path);
207        }
208    }
209
210    // 3. Global: ~/.robit/config.toml
211    let global_path = robit_home()?.join("config.toml");
212    if global_path.exists() {
213        return Ok(global_path);
214    }
215
216    Err(LlmError::ConfigError(format!(
217        "Configuration file config.toml not found.\n\
218         Please create one of the following:\n\
219         - Project-local: .robit/config.toml\n\
220         - Global: {}",
221        global_path.display()
222    )))
223}
224
225/// Resolve which model to use.
226///
227/// `default_model` uses "provider/model" format.
228/// Priority: explicit `provider_name` argument > `default_model` field > first available.
229///
230/// When `provider_name` is None, parses `default_model` (e.g. "deepseek/deepseek-chat")
231/// into provider key and model ID.
232pub fn resolve_profile(
233    config: &RobitConfig,
234    provider_name: Option<&str>,
235) -> Result<ResolvedModel, LlmError> {
236    let (provider_key, model_id) = if let Some(name) = provider_name {
237        // Explicit provider override — use its first model
238        let provider = config.providers.get(name).ok_or_else(|| {
239            LlmError::ConfigError(format!(
240                "Provider '{}' is not defined in config.toml. Available providers: {:?}",
241                name,
242                config.providers.keys().collect::<Vec<_>>()
243            ))
244        })?;
245        let first_model = provider.models.first().ok_or_else(|| {
246            LlmError::ConfigError(format!(
247                "Provider '{}' has no models defined",
248                name
249            ))
250        })?;
251        (name.to_string(), first_model.id.clone())
252    } else if let Some(ref default_model) = config.default_model {
253        parse_default_model(default_model)?
254    } else {
255        // Fall back to first available provider + first model
256        let (key, provider) = config.providers.iter().next().ok_or_else(|| {
257            LlmError::ConfigError("No providers defined in config.toml".to_string())
258        })?;
259        let first_model = provider.models.first().ok_or_else(|| {
260            LlmError::ConfigError(format!(
261                "Provider '{}' has no models defined",
262                key
263            ))
264        })?;
265        (key.clone(), first_model.id.clone())
266    };
267
268    let provider = config.providers.get(&provider_key).ok_or_else(|| {
269        LlmError::ConfigError(format!(
270            "Provider '{}' is not defined in config.toml. Available providers: {:?}",
271            provider_key,
272            config.providers.keys().collect::<Vec<_>>()
273        ))
274    })?;
275
276    // Find the matching model
277    let model = provider.models.iter().find(|m| m.id == model_id).ok_or_else(|| {
278        let available: Vec<&str> = provider.models.iter().map(|m| m.id.as_str()).collect();
279        LlmError::ConfigError(format!(
280            "Model '{}' not found in provider '{}'. Available models: {:?}",
281            model_id, provider_key, available
282        ))
283    })?;
284
285    // Validate API key
286    if provider.api_key.is_empty() || provider.api_key.starts_with("${") {
287        return Err(LlmError::ConfigError(format!(
288            "Provider '{}' API key is not configured or the environment variable is not set",
289            provider_key
290        )));
291    }
292
293    Ok(ResolvedModel {
294        profile_name: provider_key,
295        model_id: model.id.clone(),
296        base_url: provider.base_url.clone(),
297        api_key: provider.api_key.clone(),
298        max_tokens: model.max_tokens,
299        temperature: model.temperature,
300        context_window: model.context_window,
301    })
302}
303
304/// Parse "provider/model" format from default_model.
305///
306/// Returns (provider_key, model_id).
307fn parse_default_model(default_model: &str) -> Result<(String, String), LlmError> {
308    let parts: Vec<&str> = default_model.splitn(2, '/').collect();
309    if parts.len() != 2 || parts[0].is_empty() || parts[1].is_empty() {
310        return Err(LlmError::ConfigError(format!(
311            "Invalid default_model '{}' format, expected 'provider/model' (e.g. 'deepseek/deepseek-chat')",
312            default_model
313        )));
314    }
315    Ok((parts[0].to_string(), parts[1].to_string()))
316}
317
318// ============================================================================
319// Tests
320// ============================================================================
321
322#[cfg(test)]
323mod tests {
324    use super::*;
325
326    #[test]
327    fn test_resolve_env_var_with_env_set() {
328        std::env::set_var("ROBIT_TEST_KEY", "test-value-123");
329        assert_eq!(resolve_env_var("${ROBIT_TEST_KEY}"), "test-value-123");
330        std::env::remove_var("ROBIT_TEST_KEY");
331    }
332
333    #[test]
334    fn test_resolve_env_var_without_env() {
335        assert_eq!(
336            resolve_env_var("${ROBIT_NONEXISTENT_KEY}"),
337            "${ROBIT_NONEXISTENT_KEY}"
338        );
339    }
340
341    #[test]
342    fn test_resolve_env_var_plain_string() {
343        assert_eq!(resolve_env_var("plain-key"), "plain-key");
344    }
345
346    #[test]
347    fn test_parse_robit_config() {
348        let toml_str = r#"
349            default_model = "deepseek/deepseek-chat"
350
351            [providers.deepseek]
352            name = "DeepSeek"
353            base_url = "https://api.deepseek.com"
354            api_key = "sk-test-key"
355
356            [[providers.deepseek.models]]
357            id = "deepseek-chat"
358            name = "DeepSeek Chat"
359            context_window = 65536
360            max_output_tokens = 8192
361            temperature = 0.0
362            max_tokens = 4096
363
364            [[providers.deepseek.models]]
365            id = "deepseek-reasoner"
366            name = "DeepSeek Reasoner"
367            context_window = 65536
368            temperature = 0.6
369
370            [providers.qwen]
371            name = "通义千问"
372            base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
373            api_key = "sk-qwen-key"
374
375            [[providers.qwen.models]]
376            id = "qwen-max"
377            name = "Qwen Max"
378            context_window = 32768
379
380            [app]
381            log_level = "DEBUG"
382            max_steps = 10
383
384            [app.context]
385            max_output_lines = 500
386            reserve_ratio = 0.2
387
388            [app.retry]
389            max_retries = 3
390        "#;
391
392        let config: RobitConfig = toml::from_str(toml_str).unwrap();
393
394        // Default model
395        assert_eq!(config.default_model.as_deref(), Some("deepseek/deepseek-chat"));
396
397        // Providers
398        assert_eq!(config.providers.len(), 2);
399
400        // DeepSeek provider
401        let ds = &config.providers["deepseek"];
402        assert_eq!(ds.name.as_deref(), Some("DeepSeek"));
403        assert_eq!(ds.base_url, "https://api.deepseek.com");
404        assert_eq!(ds.api_key, "sk-test-key");
405        assert_eq!(ds.models.len(), 2);
406        assert_eq!(ds.models[0].id, "deepseek-chat");
407        assert_eq!(ds.models[0].context_window, Some(65536));
408        assert_eq!(ds.models[0].temperature, Some(0.0));
409        assert_eq!(ds.models[0].max_tokens, Some(4096));
410        assert_eq!(ds.models[1].id, "deepseek-reasoner");
411        assert_eq!(ds.models[1].temperature, Some(0.6));
412
413        // Qwen provider
414        let qw = &config.providers["qwen"];
415        assert_eq!(qw.name.as_deref(), Some("通义千问"));
416        assert_eq!(qw.models.len(), 1);
417        assert_eq!(qw.models[0].id, "qwen-max");
418
419        // App section
420        let app = config.app.as_ref().unwrap();
421        assert_eq!(app.log_level.as_deref(), Some("DEBUG"));
422        assert_eq!(app.max_steps, Some(10));
423        assert!(app.context.is_some());
424        assert_eq!(app.context.as_ref().unwrap().max_output_lines, Some(500));
425        assert!(app.retry.is_some());
426        assert_eq!(app.retry.as_ref().unwrap().max_retries, Some(3));
427    }
428
429    #[test]
430    fn test_parse_config_minimal() {
431        let toml_str = r#"
432            [providers.default]
433            base_url = "https://api.deepseek.com"
434            api_key = "sk-test"
435
436            [[providers.default.models]]
437            id = "deepseek-chat"
438        "#;
439
440        let config: RobitConfig = toml::from_str(toml_str).unwrap();
441        assert!(config.default_model.is_none());
442        assert!(config.app.is_none());
443        assert_eq!(config.providers.len(), 1);
444    }
445
446    #[test]
447    fn test_resolve_profile_from_default_model() {
448        let config = make_test_config();
449        let resolved = resolve_profile(&config, None).unwrap();
450        assert_eq!(resolved.profile_name, "deepseek");
451        assert_eq!(resolved.model_id, "deepseek-chat");
452        assert_eq!(resolved.base_url, "https://api.deepseek.com");
453        assert_eq!(resolved.api_key, "sk-test");
454        assert_eq!(resolved.context_window, Some(65536));
455        assert_eq!(resolved.temperature, Some(0.0));
456        assert_eq!(resolved.max_tokens, Some(4096));
457    }
458
459    #[test]
460    fn test_resolve_profile_explicit_provider() {
461        let config = make_test_config();
462        // Explicit provider — uses first model of that provider
463        let resolved = resolve_profile(&config, Some("qwen")).unwrap();
464        assert_eq!(resolved.profile_name, "qwen");
465        assert_eq!(resolved.model_id, "qwen-max");
466        assert_eq!(resolved.base_url, "https://dashscope.aliyuncs.com/compatible-mode/v1");
467    }
468
469    #[test]
470    fn test_resolve_profile_first_available() {
471        // No default_model and no explicit provider — use first available
472        let toml_str = r#"
473            [providers.deepseek]
474            base_url = "https://api.deepseek.com"
475            api_key = "sk-test"
476
477            [[providers.deepseek.models]]
478            id = "deepseek-chat"
479        "#;
480        let config: RobitConfig = toml::from_str(toml_str).unwrap();
481        let resolved = resolve_profile(&config, None).unwrap();
482        assert_eq!(resolved.profile_name, "deepseek");
483        assert_eq!(resolved.model_id, "deepseek-chat");
484    }
485
486    #[test]
487    fn test_resolve_profile_not_found() {
488        let config = make_test_config();
489        let result = resolve_profile(&config, Some("nonexistent"));
490        assert!(result.is_err());
491    }
492
493    #[test]
494    fn test_resolve_profile_model_not_found() {
495        let toml_str = r#"
496            default_model = "deepseek/nonexistent-model"
497
498            [providers.deepseek]
499            base_url = "https://api.deepseek.com"
500            api_key = "sk-test"
501
502            [[providers.deepseek.models]]
503            id = "deepseek-chat"
504        "#;
505        let config: RobitConfig = toml::from_str(toml_str).unwrap();
506        let result = resolve_profile(&config, None);
507        assert!(result.is_err());
508    }
509
510    #[test]
511    fn test_resolve_profile_invalid_default_model_format() {
512        let toml_str = r#"
513            default_model = "invalid-no-slash"
514
515            [providers.deepseek]
516            base_url = "https://api.deepseek.com"
517            api_key = "sk-test"
518
519            [[providers.deepseek.models]]
520            id = "deepseek-chat"
521        "#;
522        let config: RobitConfig = toml::from_str(toml_str).unwrap();
523        let result = resolve_profile(&config, None);
524        assert!(result.is_err());
525        assert!(result.unwrap_err().to_string().contains("Invalid default_model"));
526    }
527
528    #[test]
529    fn test_resolve_profile_empty_api_key() {
530        let toml_str = r#"
531            [providers.deepseek]
532            base_url = "https://api.deepseek.com"
533            api_key = ""
534
535            [[providers.deepseek.models]]
536            id = "deepseek-chat"
537        "#;
538        let config: RobitConfig = toml::from_str(toml_str).unwrap();
539        let result = resolve_profile(&config, None);
540        assert!(result.is_err());
541    }
542
543    #[test]
544    fn test_parse_enabled_skills() {
545        let toml_str = r#"
546            default_model = "deepseek/deepseek-chat"
547
548            [providers.deepseek]
549            base_url = "https://api.deepseek.com"
550            api_key = "sk-test"
551
552            [[providers.deepseek.models]]
553            id = "deepseek-chat"
554
555            [app]
556            enabled_skills = ["code-review", "refactor"]
557        "#;
558
559        let config: RobitConfig = toml::from_str(toml_str).unwrap();
560        let app = config.app.as_ref().unwrap();
561        assert!(app.enabled_skills.is_some());
562        let skills = app.enabled_skills.as_ref().unwrap();
563        assert_eq!(skills.len(), 2);
564        assert_eq!(skills[0], "code-review");
565        assert_eq!(skills[1], "refactor");
566    }
567
568    #[test]
569    fn test_parse_enabled_tools() {
570        let toml_str = r#"
571            default_model = "deepseek/deepseek-chat"
572
573            [providers.deepseek]
574            base_url = "https://api.deepseek.com"
575            api_key = "sk-test"
576
577            [[providers.deepseek.models]]
578            id = "deepseek-chat"
579
580            [app]
581            enabled_tools = ["read", "bash", "edit", "write", "grep", "find", "ls"]
582        "#;
583
584        let config: RobitConfig = toml::from_str(toml_str).unwrap();
585        let app = config.app.as_ref().unwrap();
586        assert!(app.enabled_tools.is_some());
587        let tools = app.enabled_tools.as_ref().unwrap();
588        assert_eq!(tools.len(), 7);
589        assert_eq!(tools[0], "read");
590        assert_eq!(tools[1], "bash");
591        assert_eq!(tools[2], "edit");
592        assert_eq!(tools[3], "write");
593        assert_eq!(tools[4], "grep");
594        assert_eq!(tools[5], "find");
595        assert_eq!(tools[6], "ls");
596    }
597
598    #[test]
599    fn test_parse_auto_approve() {
600        let toml_str = r#"
601            default_model = "deepseek/deepseek-chat"
602
603            [providers.deepseek]
604            base_url = "https://api.deepseek.com"
605            api_key = "sk-test"
606
607            [[providers.deepseek.models]]
608            id = "deepseek-chat"
609
610            [app]
611            auto_approve = true
612        "#;
613
614        let config: RobitConfig = toml::from_str(toml_str).unwrap();
615        let app = config.app.as_ref().unwrap();
616        assert_eq!(app.auto_approve, Some(true));
617    }
618
619    #[test]
620    fn test_parse_auto_approve_default_none() {
621        let toml_str = r#"
622            default_model = "deepseek/deepseek-chat"
623
624            [providers.deepseek]
625            base_url = "https://api.deepseek.com"
626            api_key = "sk-test"
627
628            [[providers.deepseek.models]]
629            id = "deepseek-chat"
630
631            [app]
632        "#;
633
634        let config: RobitConfig = toml::from_str(toml_str).unwrap();
635        let app = config.app.as_ref().unwrap();
636        assert_eq!(app.auto_approve, None);
637    }
638
639    fn make_test_config() -> RobitConfig {
640        let toml_str = r#"
641            default_model = "deepseek/deepseek-chat"
642
643            [providers.deepseek]
644            base_url = "https://api.deepseek.com"
645            api_key = "sk-test"
646
647            [[providers.deepseek.models]]
648            id = "deepseek-chat"
649            context_window = 65536
650            temperature = 0.0
651            max_tokens = 4096
652
653            [providers.qwen]
654            base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
655            api_key = "sk-qwen-test"
656
657            [[providers.qwen.models]]
658            id = "qwen-max"
659            context_window = 32768
660        "#;
661
662        toml::from_str(toml_str).unwrap()
663    }
664}
robit_ai/config.rs

robit_ai/
config.rs
