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    pub global_storage: Option<bool>,
92}
93
94#[derive(Debug, Deserialize)]
95pub struct ContextConfig {
96    pub max_output_lines: Option<usize>,
97    pub max_output_bytes: Option<usize>,
98    pub reserve_ratio: Option<f32>,
99    /// Token threshold for triggering compression (default 5000).
100    /// Only compress when removed messages exceed this token count.
101    pub compression_token_threshold: Option<usize>,
102    /// Enable/disable context compression (default true).
103    pub compression_enabled: Option<bool>,
104}
105
106#[derive(Debug, Deserialize)]
107pub struct RetryConfig {
108    pub max_retries: Option<u32>,
109    pub initial_backoff_ms: Option<u64>,
110    pub max_backoff_ms: Option<u64>,
111}
112
113// ============================================================================
114// Resolved model reference
115// ============================================================================
116
117/// A fully resolved model ready for client construction.
118///
119/// Merges provider-level settings (base_url, api_key) with model-level
120/// settings (context_window, temperature, etc).
121#[derive(Debug, Clone)]
122pub struct ResolvedModel {
123    pub profile_name: String,
124    pub model_id: String,
125    pub base_url: String,
126    pub api_key: String,
127    pub max_tokens: Option<u32>,
128    pub temperature: Option<f32>,
129    pub context_window: Option<u64>,
130}
131
132// ============================================================================
133// Loader
134// ============================================================================
135
136/// Returns the ~/.robit/ directory path.
137fn robit_home() -> Result<PathBuf, LlmError> {
138    let home = dirs::home_dir()
139        .ok_or_else(|| LlmError::ConfigError("Cannot determine home directory".to_string()))?;
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)
167        .map_err(|e| LlmError::ConfigError(format!("Failed to read {}: {}", path.display(), e)))?;
168
169    let mut config: RobitConfig = toml::from_str(&content)
170        .map_err(|e| LlmError::ConfigError(format!("Failed to parse config.toml: {}", e)))?;
171
172    // Resolve environment variables in api_key fields
173    for provider in config.providers.values_mut() {
174        provider.api_key = resolve_env_var(&provider.api_key);
175    }
176
177    Ok(config)
178}
179
180/// Load .env from ~/.robit/.env if it exists.
181pub fn load_env() {
182    if let Ok(robit_dir) = robit_home() {
183        let env_path = robit_dir.join(".env");
184        if env_path.exists() {
185            let _ = dotenvy::from_path(&env_path);
186        }
187    }
188}
189
190/// Find the config file path following the search order.
191fn find_config_path(workdir: Option<&std::path::Path>) -> Result<PathBuf, LlmError> {
192    // 1. Project-local: workdir/.robit/config.toml (if workdir provided)
193    if let Some(workdir) = workdir {
194        let local_path = workdir.join(".robit").join("config.toml");
195        if local_path.exists() {
196            return Ok(local_path);
197        }
198    }
199
200    // 2. Project-local: cwd/.robit/config.toml (if workdir not provided or no config there)
201    if let Ok(cwd) = std::env::current_dir() {
202        let local_path = cwd.join(".robit").join("config.toml");
203        if local_path.exists() {
204            return Ok(local_path);
205        }
206    }
207
208    // 3. Global: ~/.robit/config.toml
209    let global_path = robit_home()?.join("config.toml");
210    if global_path.exists() {
211        return Ok(global_path);
212    }
213
214    Err(LlmError::ConfigError(format!(
215        "Configuration file config.toml not found.\n\
216         Please create one of the following:\n\
217         - Project-local: .robit/config.toml\n\
218         - Global: {}",
219        global_path.display()
220    )))
221}
222
223/// Resolve which model to use.
224///
225/// `default_model` uses "provider/model" format.
226/// Priority: explicit `provider_name` argument > `default_model` field > first available.
227///
228/// When `provider_name` is None, parses `default_model` (e.g. "deepseek/deepseek-chat")
229/// into provider key and model ID.
230pub fn resolve_profile(
231    config: &RobitConfig,
232    provider_name: Option<&str>,
233) -> Result<ResolvedModel, LlmError> {
234    let (provider_key, model_id) = if let Some(name) = provider_name {
235        // Explicit provider override — use its first model
236        let provider = config.providers.get(name).ok_or_else(|| {
237            LlmError::ConfigError(format!(
238                "Provider '{}' is not defined in config.toml. Available providers: {:?}",
239                name,
240                config.providers.keys().collect::<Vec<_>>()
241            ))
242        })?;
243        let first_model = provider.models.first().ok_or_else(|| {
244            LlmError::ConfigError(format!("Provider '{}' has no models defined", name))
245        })?;
246        (name.to_string(), first_model.id.clone())
247    } else if let Some(ref default_model) = config.default_model {
248        parse_default_model(default_model)?
249    } else {
250        // Fall back to first available provider + first model
251        let (key, provider) = config.providers.iter().next().ok_or_else(|| {
252            LlmError::ConfigError("No providers defined in config.toml".to_string())
253        })?;
254        let first_model = provider.models.first().ok_or_else(|| {
255            LlmError::ConfigError(format!("Provider '{}' has no models defined", key))
256        })?;
257        (key.clone(), first_model.id.clone())
258    };
259
260    let provider = config.providers.get(&provider_key).ok_or_else(|| {
261        LlmError::ConfigError(format!(
262            "Provider '{}' is not defined in config.toml. Available providers: {:?}",
263            provider_key,
264            config.providers.keys().collect::<Vec<_>>()
265        ))
266    })?;
267
268    // Find the matching model
269    let model = provider
270        .models
271        .iter()
272        .find(|m| m.id == model_id)
273        .ok_or_else(|| {
274            let available: Vec<&str> = provider.models.iter().map(|m| m.id.as_str()).collect();
275            LlmError::ConfigError(format!(
276                "Model '{}' not found in provider '{}'. Available models: {:?}",
277                model_id, provider_key, available
278            ))
279        })?;
280
281    // Validate API key
282    if provider.api_key.is_empty() || provider.api_key.starts_with("${") {
283        return Err(LlmError::ConfigError(format!(
284            "Provider '{}' API key is not configured or the environment variable is not set",
285            provider_key
286        )));
287    }
288
289    Ok(ResolvedModel {
290        profile_name: provider_key,
291        model_id: model.id.clone(),
292        base_url: provider.base_url.clone(),
293        api_key: provider.api_key.clone(),
294        max_tokens: model.max_tokens,
295        temperature: model.temperature,
296        context_window: model.context_window,
297    })
298}
299
300/// Parse "provider/model" format from default_model.
301///
302/// Returns (provider_key, model_id).
303fn parse_default_model(default_model: &str) -> Result<(String, String), LlmError> {
304    let parts: Vec<&str> = default_model.splitn(2, '/').collect();
305    if parts.len() != 2 || parts[0].is_empty() || parts[1].is_empty() {
306        return Err(LlmError::ConfigError(format!(
307            "Invalid default_model '{}' format, expected 'provider/model' (e.g. 'deepseek/deepseek-chat')",
308            default_model
309        )));
310    }
311    Ok((parts[0].to_string(), parts[1].to_string()))
312}
313
314// ============================================================================
315// Tests
316// ============================================================================
317
318#[cfg(test)]
319mod tests {
320    use super::*;
321
322    #[test]
323    fn test_resolve_env_var_with_env_set() {
324        std::env::set_var("ROBIT_TEST_KEY", "test-value-123");
325        assert_eq!(resolve_env_var("${ROBIT_TEST_KEY}"), "test-value-123");
326        std::env::remove_var("ROBIT_TEST_KEY");
327    }
328
329    #[test]
330    fn test_resolve_env_var_without_env() {
331        assert_eq!(
332            resolve_env_var("${ROBIT_NONEXISTENT_KEY}"),
333            "${ROBIT_NONEXISTENT_KEY}"
334        );
335    }
336
337    #[test]
338    fn test_resolve_env_var_plain_string() {
339        assert_eq!(resolve_env_var("plain-key"), "plain-key");
340    }
341
342    #[test]
343    fn test_parse_robit_config() {
344        let toml_str = r#"
345            default_model = "deepseek/deepseek-chat"
346
347            [providers.deepseek]
348            name = "DeepSeek"
349            base_url = "https://api.deepseek.com"
350            api_key = "sk-test-key"
351
352            [[providers.deepseek.models]]
353            id = "deepseek-chat"
354            name = "DeepSeek Chat"
355            context_window = 65536
356            max_output_tokens = 8192
357            temperature = 0.0
358            max_tokens = 4096
359
360            [[providers.deepseek.models]]
361            id = "deepseek-reasoner"
362            name = "DeepSeek Reasoner"
363            context_window = 65536
364            temperature = 0.6
365
366            [providers.qwen]
367            name = "通义千问"
368            base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
369            api_key = "sk-qwen-key"
370
371            [[providers.qwen.models]]
372            id = "qwen-max"
373            name = "Qwen Max"
374            context_window = 32768
375
376            [app]
377            log_level = "DEBUG"
378            max_steps = 10
379            global_storage = true
380
381            [app.context]
382            max_output_lines = 500
383            reserve_ratio = 0.2
384
385            [app.retry]
386            max_retries = 3
387        "#;
388
389        let config: RobitConfig = toml::from_str(toml_str).unwrap();
390
391        // Default model
392        assert_eq!(
393            config.default_model.as_deref(),
394            Some("deepseek/deepseek-chat")
395        );
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_eq!(app.global_storage, Some(true));
424        assert!(app.context.is_some());
425        assert_eq!(app.context.as_ref().unwrap().max_output_lines, Some(500));
426        assert!(app.retry.is_some());
427        assert_eq!(app.retry.as_ref().unwrap().max_retries, Some(3));
428    }
429
430    #[test]
431    fn test_parse_config_minimal() {
432        let toml_str = r#"
433            [providers.default]
434            base_url = "https://api.deepseek.com"
435            api_key = "sk-test"
436
437            [[providers.default.models]]
438            id = "deepseek-chat"
439        "#;
440
441        let config: RobitConfig = toml::from_str(toml_str).unwrap();
442        assert!(config.default_model.is_none());
443        assert!(config.app.is_none());
444        assert_eq!(config.providers.len(), 1);
445    }
446
447    #[test]
448    fn test_resolve_profile_from_default_model() {
449        let config = make_test_config();
450        let resolved = resolve_profile(&config, None).unwrap();
451        assert_eq!(resolved.profile_name, "deepseek");
452        assert_eq!(resolved.model_id, "deepseek-chat");
453        assert_eq!(resolved.base_url, "https://api.deepseek.com");
454        assert_eq!(resolved.api_key, "sk-test");
455        assert_eq!(resolved.context_window, Some(65536));
456        assert_eq!(resolved.temperature, Some(0.0));
457        assert_eq!(resolved.max_tokens, Some(4096));
458    }
459
460    #[test]
461    fn test_resolve_profile_explicit_provider() {
462        let config = make_test_config();
463        // Explicit provider — uses first model of that provider
464        let resolved = resolve_profile(&config, Some("qwen")).unwrap();
465        assert_eq!(resolved.profile_name, "qwen");
466        assert_eq!(resolved.model_id, "qwen-max");
467        assert_eq!(
468            resolved.base_url,
469            "https://dashscope.aliyuncs.com/compatible-mode/v1"
470        );
471    }
472
473    #[test]
474    fn test_resolve_profile_first_available() {
475        // No default_model and no explicit provider — use first available
476        let toml_str = r#"
477            [providers.deepseek]
478            base_url = "https://api.deepseek.com"
479            api_key = "sk-test"
480
481            [[providers.deepseek.models]]
482            id = "deepseek-chat"
483        "#;
484        let config: RobitConfig = toml::from_str(toml_str).unwrap();
485        let resolved = resolve_profile(&config, None).unwrap();
486        assert_eq!(resolved.profile_name, "deepseek");
487        assert_eq!(resolved.model_id, "deepseek-chat");
488    }
489
490    #[test]
491    fn test_resolve_profile_not_found() {
492        let config = make_test_config();
493        let result = resolve_profile(&config, Some("nonexistent"));
494        assert!(result.is_err());
495    }
496
497    #[test]
498    fn test_resolve_profile_model_not_found() {
499        let toml_str = r#"
500            default_model = "deepseek/nonexistent-model"
501
502            [providers.deepseek]
503            base_url = "https://api.deepseek.com"
504            api_key = "sk-test"
505
506            [[providers.deepseek.models]]
507            id = "deepseek-chat"
508        "#;
509        let config: RobitConfig = toml::from_str(toml_str).unwrap();
510        let result = resolve_profile(&config, None);
511        assert!(result.is_err());
512    }
513
514    #[test]
515    fn test_resolve_profile_invalid_default_model_format() {
516        let toml_str = r#"
517            default_model = "invalid-no-slash"
518
519            [providers.deepseek]
520            base_url = "https://api.deepseek.com"
521            api_key = "sk-test"
522
523            [[providers.deepseek.models]]
524            id = "deepseek-chat"
525        "#;
526        let config: RobitConfig = toml::from_str(toml_str).unwrap();
527        let result = resolve_profile(&config, None);
528        assert!(result.is_err());
529        assert!(result
530            .unwrap_err()
531            .to_string()
532            .contains("Invalid default_model"));
533    }
534
535    #[test]
536    fn test_resolve_profile_empty_api_key() {
537        let toml_str = r#"
538            [providers.deepseek]
539            base_url = "https://api.deepseek.com"
540            api_key = ""
541
542            [[providers.deepseek.models]]
543            id = "deepseek-chat"
544        "#;
545        let config: RobitConfig = toml::from_str(toml_str).unwrap();
546        let result = resolve_profile(&config, None);
547        assert!(result.is_err());
548    }
549
550    #[test]
551    fn test_parse_enabled_skills() {
552        let toml_str = r#"
553            default_model = "deepseek/deepseek-chat"
554
555            [providers.deepseek]
556            base_url = "https://api.deepseek.com"
557            api_key = "sk-test"
558
559            [[providers.deepseek.models]]
560            id = "deepseek-chat"
561
562            [app]
563            enabled_skills = ["code-review", "refactor"]
564        "#;
565
566        let config: RobitConfig = toml::from_str(toml_str).unwrap();
567        let app = config.app.as_ref().unwrap();
568        assert!(app.enabled_skills.is_some());
569        let skills = app.enabled_skills.as_ref().unwrap();
570        assert_eq!(skills.len(), 2);
571        assert_eq!(skills[0], "code-review");
572        assert_eq!(skills[1], "refactor");
573    }
574
575    #[test]
576    fn test_parse_enabled_tools() {
577        let toml_str = r#"
578            default_model = "deepseek/deepseek-chat"
579
580            [providers.deepseek]
581            base_url = "https://api.deepseek.com"
582            api_key = "sk-test"
583
584            [[providers.deepseek.models]]
585            id = "deepseek-chat"
586
587            [app]
588            enabled_tools = ["read", "bash", "edit", "write", "grep", "find", "ls"]
589        "#;
590
591        let config: RobitConfig = toml::from_str(toml_str).unwrap();
592        let app = config.app.as_ref().unwrap();
593        assert!(app.enabled_tools.is_some());
594        let tools = app.enabled_tools.as_ref().unwrap();
595        assert_eq!(tools.len(), 7);
596        assert_eq!(tools[0], "read");
597        assert_eq!(tools[1], "bash");
598        assert_eq!(tools[2], "edit");
599        assert_eq!(tools[3], "write");
600        assert_eq!(tools[4], "grep");
601        assert_eq!(tools[5], "find");
602        assert_eq!(tools[6], "ls");
603    }
604
605    #[test]
606    fn test_parse_auto_approve() {
607        let toml_str = r#"
608            default_model = "deepseek/deepseek-chat"
609
610            [providers.deepseek]
611            base_url = "https://api.deepseek.com"
612            api_key = "sk-test"
613
614            [[providers.deepseek.models]]
615            id = "deepseek-chat"
616
617            [app]
618            auto_approve = true
619        "#;
620
621        let config: RobitConfig = toml::from_str(toml_str).unwrap();
622        let app = config.app.as_ref().unwrap();
623        assert_eq!(app.auto_approve, Some(true));
624    }
625
626    #[test]
627    fn test_parse_auto_approve_default_none() {
628        let toml_str = r#"
629            default_model = "deepseek/deepseek-chat"
630
631            [providers.deepseek]
632            base_url = "https://api.deepseek.com"
633            api_key = "sk-test"
634
635            [[providers.deepseek.models]]
636            id = "deepseek-chat"
637
638            [app]
639        "#;
640
641        let config: RobitConfig = toml::from_str(toml_str).unwrap();
642        let app = config.app.as_ref().unwrap();
643        assert_eq!(app.auto_approve, None);
644    }
645
646    fn make_test_config() -> RobitConfig {
647        let toml_str = r#"
648            default_model = "deepseek/deepseek-chat"
649
650            [providers.deepseek]
651            base_url = "https://api.deepseek.com"
652            api_key = "sk-test"
653
654            [[providers.deepseek.models]]
655            id = "deepseek-chat"
656            context_window = 65536
657            temperature = 0.0
658            max_tokens = 4096
659
660            [providers.qwen]
661            base_url = "https://dashscope.aliyuncs.com/compatible-mode/v1"
662            api_key = "sk-qwen-test"
663
664            [[providers.qwen.models]]
665            id = "qwen-max"
666            context_window = 32768
667        "#;
668
669        toml::from_str(toml_str).unwrap()
670    }
671}
robit_ai/config.rs

robit_ai/
config.rs
