cookie_agent_config/
loader.rs

1use std::{
2    collections::BTreeMap,
3    path::{Path, PathBuf},
4    sync::Arc,
5};
6
7use cookie_agent_identity::AgentId;
8use cookie_agent_identity::ProviderId;
9use cookie_agent_models::ProviderDefinition;
10use cookie_agent_protocol::paths;
11use indexmap::IndexMap;
12use serde::Deserialize as _;
13
14use crate::{
15    AgentDocument, AgentDocumentSource, AgentMdConfig, AgentRegistry, ApprovalConfig, ConfigError,
16    ContextCompactionConfig, DelegationConfig, EngineConfig, MessagingConfig, ModelRetryConfig,
17    ServerConfig, SessionTitleConfig, ToolOutputConfig,
18    agent_document::parse_agent,
19    runtime::{RawRuntimeLayer, apply_settings, validate_runtime},
20    secure_fs::{
21        LayerRoot, open_layer_root, open_optional_directory, read_optional_file,
22        read_required_file, regular_file_at,
23    },
24    toml_values::{
25        SensitiveJsonValue, SensitiveProviderValues, SensitiveTomlValue,
26        interpolate_provider_values, safe_toml_error, validate_toml_value, zeroize_toml_value,
27    },
28};
29
30const MAX_CONFIG_BYTES: u64 = 1024 * 1024;
31const MAX_AGENT_BYTES: u64 = 256 * 1024;
32const BUILT_IN_CONFIG: &str = include_str!("defaults.toml");
33
34#[derive(Clone, Debug)]
35pub struct LoadedConfiguration {
36    /// The settings tables of the authored `config.toml`; `[mcp]` and
37    /// `[plugins]` decode onto the sibling fields instead. The field name
38    /// reflects that the config file *is* the runtime settings. Keys are
39    /// top-level, so `[messaging]` maps to `runtime.messaging`, not
40    /// `[runtime.messaging]`.
41    pub runtime: EngineConfig,
42    pub agents: BTreeMap<AgentId, AgentDocument>,
43    pub agent_presets: BTreeMap<String, BTreeMap<AgentId, AgentDocument>>,
44    pub mcp_servers: BTreeMap<String, LoadedMcpServer>,
45    pub user_mcp_servers: BTreeMap<String, crate::McpServerConfig>,
46    pub workspace_mcp_servers: BTreeMap<String, crate::McpServerConfig>,
47    pub plugins: IndexMap<String, crate::PluginConfig>,
48    pub config_paths: ConfigLayerPaths,
49    pub skills: crate::SkillRegistry,
50}
51
52#[derive(Clone, Debug, Default)]
53pub struct ConfigLayerPaths {
54    pub user: Option<PathBuf>,
55    pub workspace: Option<PathBuf>,
56}
57
58#[derive(Clone, Copy, Debug, Eq, PartialEq)]
59pub enum McpServerSource {
60    UserFile,
61    WorkspaceFile,
62    Runtime,
63}
64
65#[derive(Clone, Debug)]
66pub struct LoadedMcpServer {
67    pub source: McpServerSource,
68    pub config: crate::McpServerConfig,
69}
70
71impl LoadedConfiguration {
72    #[must_use]
73    pub fn agent_registry(&self) -> AgentRegistry {
74        AgentRegistry::from_validated(self.agents.clone())
75    }
76
77    #[must_use]
78    pub fn agent_preset_registries(&self) -> BTreeMap<String, AgentRegistry> {
79        self.agent_presets
80            .iter()
81            .map(|(name, agents)| (name.clone(), AgentRegistry::from_validated(agents.clone())))
82            .collect()
83    }
84}
85
86/// Loads user config and exact-cwd `.cookie-agent` with no upward search.
87pub fn load(cwd: &Path) -> Result<LoadedConfiguration, ConfigError> {
88    let cwd = cwd.canonicalize().map_err(ConfigError::Io)?;
89    let user = user_root();
90    let mut loaded = load_from_roots(user.as_deref(), Some(&cwd.join(".cookie-agent")))?;
91    loaded.skills = crate::load_skills(&cwd)?;
92    Ok(loaded)
93}
94
95/// Loads explicit roots, primarily for composition and adversarial tests.
96pub fn load_from_roots(
97    user_root: Option<&Path>,
98    workspace_root: Option<&Path>,
99) -> Result<LoadedConfiguration, ConfigError> {
100    let user = user_root
101        .map(|path| open_layer_root(path, AgentDocumentSource::User))
102        .transpose()?
103        .flatten();
104    let workspace = workspace_root
105        .map(|path| open_layer_root(path, AgentDocumentSource::Workspace))
106        .transpose()?
107        .flatten();
108    let mut runtime = EngineConfig {
109        server: ServerConfig::default(),
110        tool_output: ToolOutputConfig::default(),
111        agent_md: AgentMdConfig::default(),
112        approval: ApprovalConfig::default(),
113        model_retry: ModelRetryConfig::default(),
114        context_compaction: ContextCompactionConfig::default(),
115        session_title: SessionTitleConfig::default(),
116        delegation: DelegationConfig::default(),
117        messaging: MessagingConfig::default(),
118        pricing: Default::default(),
119        headers: BTreeMap::new(),
120        providers: BTreeMap::new(),
121    };
122    let mut agents = BTreeMap::new();
123    let mut preset_agents = BTreeMap::<String, BTreeMap<AgentId, AgentDocument>>::new();
124    let mut mcp_servers = BTreeMap::new();
125    let mut user_mcp_servers = BTreeMap::new();
126    let mut workspace_mcp_servers = BTreeMap::new();
127    let mut plugins = IndexMap::new();
128    let built_in_text =
129        BUILT_IN_CONFIG.replace("${cookie_agent_version}", env!("CARGO_PKG_VERSION"));
130    let built_in_path = Path::new("<built-in>/config.toml");
131    let mut built_in_value = built_in_text.parse::<toml::Value>().map_err(|error| {
132        ConfigError::Toml(format!(
133            "{}: {}",
134            built_in_path.display(),
135            safe_toml_error(&built_in_text, &error)
136        ))
137    })?;
138    interpolate_provider_values(&mut built_in_value, &mut Vec::new())?;
139    let built_in_layer = decode_runtime_layer(&mut built_in_value, &built_in_text, built_in_path)?;
140    apply_settings(&mut runtime, &built_in_layer);
141    for root in [user.as_ref(), workspace.as_ref()].into_iter().flatten() {
142        if let Some((mut layer, source_text)) = root.load_runtime()? {
143            apply_settings(&mut runtime, &layer);
144            if let Some(mcp) = layer.mcp.take() {
145                let source = match root.source {
146                    AgentDocumentSource::User => McpServerSource::UserFile,
147                    AgentDocumentSource::Workspace => McpServerSource::WorkspaceFile,
148                    AgentDocumentSource::BuiltIn => McpServerSource::Runtime,
149                };
150                for (name, config) in mcp.servers {
151                    match source {
152                        McpServerSource::UserFile => {
153                            user_mcp_servers.insert(name.clone(), config.clone());
154                        }
155                        McpServerSource::WorkspaceFile => {
156                            workspace_mcp_servers.insert(name.clone(), config.clone());
157                        }
158                        McpServerSource::Runtime => unreachable!("file configuration layer"),
159                    }
160                    mcp_servers.insert(name, LoadedMcpServer { source, config });
161                }
162            }
163            if let Some(layer_plugins) = layer.plugins.take() {
164                plugins.extend(plugins_in_authored_order(layer_plugins, &source_text).map(
165                    |(name, config)| {
166                        (
167                            name,
168                            (
169                                config,
170                                root.path.join("config.toml"),
171                                Arc::clone(&source_text),
172                            ),
173                        )
174                    },
175                ));
176            }
177            for (id, mut value) in std::mem::take(&mut layer.providers) {
178                let json = SensitiveJsonValue::from_toml(value.take());
179                let provider = ProviderDefinition::deserialize(json.value()).map_err(|error| {
180                    config_decode_error(
181                        &root.path.join("config.toml"),
182                        None,
183                        &format!("provider `{id}`: {error}"),
184                    )
185                })?;
186                runtime.providers.insert(id, provider);
187            }
188        }
189        let layer_agents = root.load_agents()?;
190        for (id, document) in layer_agents.shared {
191            agents.insert(id, document);
192        }
193        for (preset, documents) in layer_agents.presets {
194            preset_agents.entry(preset).or_default().extend(documents);
195        }
196    }
197    AgentRegistry::validate_ref(&agents)?;
198    let agent_presets = preset_agents
199        .into_iter()
200        .map(|(preset, documents)| {
201            let mut effective = agents.clone();
202            effective.extend(documents);
203            AgentRegistry::validate_ref(&effective)?;
204            Ok((preset, effective))
205        })
206        .collect::<Result<_, ConfigError>>()?;
207    runtime.pricing.models = runtime
208        .providers
209        .iter()
210        .flat_map(|(provider_id, provider)| {
211            let prices: Vec<_> = match provider {
212                ProviderDefinition::ModelsDev(provider) => provider
213                    .model_overrides
214                    .iter()
215                    .filter_map(|(id, model)| model.pricing.map(|pricing| (id, pricing)))
216                    .collect(),
217                ProviderDefinition::Custom(provider) => provider
218                    .models
219                    .iter()
220                    .filter_map(|(id, model)| model.pricing.map(|pricing| (id, pricing)))
221                    .collect(),
222            };
223            prices.into_iter().map(|(id, pricing)| {
224                (
225                    cookie_agent_identity::ModelKey::new(provider_id.clone(), id.clone())
226                        .expect("validated model IDs"),
227                    pricing,
228                )
229            })
230        })
231        .collect();
232    validate_runtime(&runtime)?;
233    for (name, server) in &mcp_servers {
234        server.config.validate(name)?;
235    }
236    for (name, (plugin, path, text)) in &plugins {
237        if let Some(field) = plugin.invalid_field(name) {
238            return Err(config_decode_error(
239                path,
240                Some(text),
241                &format!("plugin `{name}`: invalid field `{field}`"),
242            ));
243        }
244    }
245    let plugins = plugins
246        .into_iter()
247        .map(|(name, (config, _, _))| (name, config))
248        .collect();
249    Ok(LoadedConfiguration {
250        runtime,
251        agents,
252        agent_presets,
253        mcp_servers,
254        user_mcp_servers,
255        workspace_mcp_servers,
256        plugins,
257        config_paths: ConfigLayerPaths {
258            user: user_root.map(|root| root.join("config.toml")),
259            workspace: workspace_root.map(|root| root.join("config.toml")),
260        },
261        skills: crate::SkillRegistry::default(),
262    })
263}
264
265fn plugins_in_authored_order(
266    mut plugins: crate::PluginsConfig,
267    source_text: &str,
268) -> impl Iterator<Item = (String, crate::PluginConfig)> {
269    let document = source_text
270        .parse::<toml_edit::DocumentMut>()
271        .expect("validated TOML reparses for plugin ordering");
272    let names: Vec<String> = document
273        .get("plugins")
274        .and_then(toml_edit::Item::as_table)
275        .map(|table| table.iter().map(|(name, _)| name.to_owned()).collect())
276        .unwrap_or_default();
277    names.into_iter().filter_map(move |name| {
278        plugins
279            .plugins
280            .shift_remove(&name)
281            .map(|config| (name, config))
282    })
283}
284
285impl LayerRoot {
286    fn load_runtime(&self) -> Result<Option<(RawRuntimeLayer, Arc<str>)>, ConfigError> {
287        let Some(bytes) = read_optional_file(&self.path, "config.toml", MAX_CONFIG_BYTES)? else {
288            return Ok(None);
289        };
290        let path = self.path.join("config.toml");
291        let text = std::str::from_utf8(&bytes)
292            .map_err(|_| ConfigError::Toml(format!("{}: content is not UTF-8", path.display())))?;
293        let value = text.parse::<toml::Value>().map_err(|error| {
294            ConfigError::Toml(format!(
295                "{}: {}",
296                path.display(),
297                safe_toml_error(text, &error)
298            ))
299        })?;
300        let mut value = SensitiveTomlValue::new(value);
301        validate_toml_value(value.value()).map_err(|_| {
302            ConfigError::Toml(format!(
303                "{}: TOML resource limit or unsupported datetime exceeded",
304                path.display()
305            ))
306        })?;
307        interpolate_provider_values(value.value_mut(), &mut Vec::new())?;
308        let layer = decode_runtime_layer(value.value_mut(), text, &path)?;
309        Ok(Some((layer, Arc::from(text))))
310    }
311
312    fn load_agents(&self) -> Result<AgentLayer, ConfigError> {
313        let Some(directory) = open_optional_directory(&self.path, "agents")? else {
314            return Ok(AgentLayer::default());
315        };
316        let names = sorted_entry_names(&directory)?;
317        let mut layer = AgentLayer::default();
318        for name in names {
319            if name.ends_with(".md") {
320                load_agent_document(&directory, &name, self.source, &mut layer.shared)?;
321                continue;
322            }
323            if regular_file_at(&directory, &name)? {
324                continue;
325            }
326            let preset_directory =
327                open_optional_directory(&directory, &name)?.ok_or(ConfigError::UnsafePath)?;
328            AgentId::new(&name).map_err(|_| ConfigError::AgentPresetName {
329                path: preset_directory.clone(),
330                name: name.clone(),
331            })?;
332            let mut documents = BTreeMap::new();
333            for document_name in sorted_entry_names(&preset_directory)? {
334                if !document_name.ends_with(".md") {
335                    return Err(ConfigError::UnsafePath);
336                }
337                load_agent_document(
338                    &preset_directory,
339                    &document_name,
340                    self.source,
341                    &mut documents,
342                )?;
343            }
344            layer.presets.insert(name, documents);
345        }
346        Ok(layer)
347    }
348}
349
350#[derive(Default)]
351struct AgentLayer {
352    shared: BTreeMap<AgentId, AgentDocument>,
353    presets: BTreeMap<String, BTreeMap<AgentId, AgentDocument>>,
354}
355
356fn sorted_entry_names(directory: &Path) -> Result<Vec<String>, ConfigError> {
357    let mut names = std::fs::read_dir(directory)
358        .map_err(ConfigError::Io)?
359        .map(|entry| {
360            entry
361                .map_err(ConfigError::Io)?
362                .file_name()
363                .into_string()
364                .map_err(|_| ConfigError::UnsafePath)
365        })
366        .collect::<Result<Vec<_>, _>>()?;
367    names.sort_by(|left, right| left.as_bytes().cmp(right.as_bytes()));
368    Ok(names)
369}
370
371fn load_agent_document(
372    directory: &Path,
373    name: &str,
374    source: AgentDocumentSource,
375    documents: &mut BTreeMap<AgentId, AgentDocument>,
376) -> Result<(), ConfigError> {
377    if !regular_file_at(directory, name)? {
378        return Err(ConfigError::UnsafePath);
379    }
380    let stem = name.strip_suffix(".md").expect("suffix checked");
381    let id = AgentId::new(stem).map_err(|_| ConfigError::AgentFilename(name.to_owned()))?;
382    let bytes = read_required_file(directory, name, MAX_AGENT_BYTES)?;
383    let document = parse_agent(id.clone(), &bytes, source, &directory.join(name))?;
384    if documents.insert(id.clone(), document).is_some() {
385        return Err(ConfigError::DuplicateAgent(id));
386    }
387    Ok(())
388}
389
390fn decode_runtime_layer(
391    value: &mut toml::Value,
392    text: &str,
393    path: &Path,
394) -> Result<RawRuntimeLayer, ConfigError> {
395    const ALLOWED: &[&str] = &[
396        "server",
397        "tool_output",
398        "agent_md",
399        "approval",
400        "model_retry",
401        "context_compaction",
402        "session_title",
403        "delegation",
404        "messaging",
405        "headers",
406        "mcp",
407        "plugins",
408        "providers",
409    ];
410    let table = value.as_table_mut().ok_or_else(|| {
411        ConfigError::Toml(format!("{}: top level must be a table", path.display()))
412    })?;
413    if table.contains_key("schema_version") {
414        return Err(ConfigError::ConfigSchemaRemoved {
415            path: path.to_owned(),
416            line: key_line(text, "schema_version").unwrap_or(1),
417        });
418    }
419    if let Some(key) = table.keys().find(|key| !ALLOWED.contains(&key.as_str())) {
420        return Err(ConfigError::Toml(format!(
421            "{} line {}: unknown top-level field `{key}`; remove it",
422            path.display(),
423            key_line(text, key).unwrap_or(1)
424        )));
425    }
426    if table
427        .get("context_compaction")
428        .and_then(toml::Value::as_table)
429        .is_some_and(|context| {
430            context.contains_key("trigger") && context.contains_key("buffer_tokens")
431        })
432    {
433        return Err(ConfigError::Toml(format!(
434            "{}: context_compaction.trigger and context_compaction.buffer_tokens cannot both be set",
435            path.display()
436        )));
437    }
438    let providers = match table.remove("providers") {
439        None => SensitiveProviderValues::new(),
440        Some(toml::Value::Table(values)) => values
441            .into_iter()
442            .map(|(id, value)| {
443                let value = SensitiveTomlValue::new(value);
444                ProviderId::new(id).map(|id| (id, value)).map_err(|_| {
445                    ConfigError::Toml(format!("{}: invalid provider ID", path.display()))
446                })
447            })
448            .collect::<Result<_, _>>()?,
449        Some(mut other) => {
450            zeroize_toml_value(&mut other);
451            return Err(ConfigError::Toml(format!(
452                "{}: `providers` must be a table",
453                path.display()
454            )));
455        }
456    };
457    let owned = std::mem::replace(value, toml::Value::Table(Default::default()));
458    let json = SensitiveJsonValue::from_toml(owned);
459    let mut layer = RawRuntimeLayer::deserialize(json.value())
460        .map_err(|error| config_decode_error(path, Some(text), &error.to_string()))?;
461    layer.providers = providers;
462    Ok(layer)
463}
464
465fn config_decode_error(path: &Path, text: Option<&str>, message: &str) -> ConfigError {
466    let field = extract_serde_field(message, "unknown field")
467        .map(|field| (field, "unknown field"))
468        .or_else(|| {
469            extract_serde_field(message, "missing field").map(|field| (field, "missing field"))
470        })
471        .or_else(|| {
472            extract_serde_field(message, "invalid field").map(|field| (field, "invalid field"))
473        });
474    let detail = field.as_ref().map_or_else(
475        || "malformed configuration content".to_owned(),
476        |(field, kind)| {
477            let line = text.and_then(|text| key_line(text, field));
478            let advice = match *kind {
479                "unknown field" => "remove it",
480                "missing field" => "add it",
481                "invalid field" => "fix it",
482                _ => unreachable!("known configuration diagnostic kind"),
483            };
484            match line {
485                Some(line) => format!("line {line}: {kind} `{field}`; {advice}"),
486                None => format!("{kind} `{field}`; {advice}"),
487            }
488        },
489    );
490    ConfigError::Toml(format!("{}: {detail}", path.display()))
491}
492
493fn extract_serde_field(message: &str, marker: &str) -> Option<String> {
494    let rest = message.split_once(marker)?.1.trim_start();
495    let quote = rest.chars().next()?;
496    if !matches!(quote, '`' | '\'' | '"') {
497        return None;
498    }
499    rest[quote.len_utf8()..]
500        .split_once(quote)
501        .map(|(field, _)| field.to_owned())
502}
503
504fn key_line(text: &str, key: &str) -> Option<usize> {
505    text.lines()
506        .position(|line| {
507            line.split_once('=')
508                .is_some_and(|(candidate, _)| candidate.trim().trim_matches('"') == key)
509        })
510        .map(|index| index + 1)
511}
512
513fn user_root() -> Option<PathBuf> {
514    paths::user_data_root().ok()
515}
516
517#[cfg(test)]
518mod agent_loader_tests {
519    use super::*;
520
521    #[test]
522    fn duplicate_agent_ids_in_one_preset_map_are_rejected() {
523        let directory = tempfile::tempdir().unwrap();
524        std::fs::write(
525            directory.path().join("worker.md"),
526            "---\ndescription: Worker\nmode: primary\nenabled: true\nmodels: [{ model: \"custom.test/model\" }]\npermissions: {}\n---\nPrompt.\n",
527        )
528        .unwrap();
529        let mut documents = BTreeMap::new();
530        load_agent_document(
531            directory.path(),
532            "worker.md",
533            AgentDocumentSource::Workspace,
534            &mut documents,
535        )
536        .unwrap();
537        assert!(matches!(
538            load_agent_document(
539                directory.path(),
540                "worker.md",
541                AgentDocumentSource::Workspace,
542                &mut documents,
543            ),
544            Err(ConfigError::DuplicateAgent(id)) if id.as_str() == "worker"
545        ));
546    }
547}
cookie_agent_config/loader.rs

cookie_agent_config/
loader.rs
