cookie_agent_config/
mcp_store.rs

1use std::{
2    fs::{self, OpenOptions},
3    io::Write as _,
4    path::Path,
5};
6
7use serde::Serialize;
8use toml_edit::{DocumentMut, Item, Table};
9
10use crate::{ConfigError, McpServerConfig, load_from_roots};
11
12#[derive(Serialize)]
13struct ServerDocument<'a> {
14    mcp: McpDocument<'a>,
15}
16
17#[derive(Serialize)]
18struct McpDocument<'a> {
19    servers: std::collections::BTreeMap<&'a str, &'a McpServerConfig>,
20}
21
22/// Replace one `[mcp.servers.<name>]` table while preserving every unrelated
23/// item and comment. The candidate is accepted only when the strict loader can
24/// parse it as a complete configuration layer.
25pub fn write_mcp_server(
26    path: &Path,
27    name: &str,
28    config: &McpServerConfig,
29) -> Result<(), ConfigError> {
30    write_mcp_server_observed(path, name, config, |_| {})
31}
32
33#[derive(Clone, Copy, Debug, Eq, PartialEq)]
34enum CommitStep {
35    TemporarySynced,
36    IdentityChecked,
37    Renamed,
38    ParentSynced,
39}
40
41fn write_mcp_server_observed(
42    path: &Path,
43    name: &str,
44    config: &McpServerConfig,
45    mut observe: impl FnMut(CommitStep),
46) -> Result<(), ConfigError> {
47    config.validate(name)?;
48    let (original, originally_existed) = match fs::read_to_string(path) {
49        Ok(text) => (text, true),
50        Err(error) if error.kind() == std::io::ErrorKind::NotFound => (String::new(), false),
51        Err(error) => return Err(ConfigError::Io(error)),
52    };
53    let mut document = original
54        .parse::<DocumentMut>()
55        .map_err(|error| ConfigError::Toml(format!("{}: {error}", path.display())))?;
56    let encoded = toml::to_string(&ServerDocument {
57        mcp: McpDocument {
58            servers: std::collections::BTreeMap::from([(name, config)]),
59        },
60    })
61    .map_err(|error| ConfigError::Toml(format!("{}: {error}", path.display())))?;
62    let encoded = encoded
63        .parse::<DocumentMut>()
64        .map_err(|error| ConfigError::Toml(format!("{}: {error}", path.display())))?;
65    let replacement = encoded["mcp"]["servers"][name].clone();
66
67    if !document.as_table().contains_key("mcp") {
68        document["mcp"] = Item::Table(Table::new());
69    }
70    if !document["mcp"]
71        .as_table()
72        .is_some_and(|table| table.contains_key("servers"))
73    {
74        document["mcp"]["servers"] = Item::Table(Table::new());
75    }
76    document["mcp"]["servers"][name] = replacement;
77    let candidate = document.to_string();
78
79    let validation_root = tempfile_root(path, &candidate)?;
80    let validation = load_from_roots(Some(&validation_root), None);
81    let _ = fs::remove_dir_all(&validation_root);
82    validation?;
83
84    let parent = path.parent().ok_or(ConfigError::UnsafePath)?;
85    fs::create_dir_all(parent).map_err(ConfigError::Io)?;
86    #[cfg(unix)]
87    let temporary = parent.join(format!(
88        ".{}.{}.tmp",
89        path.file_name()
90            .and_then(|name| name.to_str())
91            .unwrap_or("config.toml"),
92        std::process::id()
93    ));
94    #[cfg(unix)]
95    let mut open_options = OpenOptions::new();
96    #[cfg(unix)]
97    open_options.write(true).create(true).truncate(true);
98    #[cfg(unix)]
99    {
100        use std::os::unix::fs::OpenOptionsExt;
101        open_options.mode(0o600);
102    }
103    #[cfg(unix)]
104    let mut temporary_file = open_options.open(&temporary).map_err(ConfigError::Io)?;
105    #[cfg(windows)]
106    let (temporary, mut temporary_file) = create_windows_temporary(parent, path)?;
107    #[cfg(unix)]
108    {
109        // A stale temp file from a crashed write keeps its old mode; force it.
110        use std::os::unix::fs::PermissionsExt;
111        if let Err(error) = temporary_file.set_permissions(fs::Permissions::from_mode(0o600)) {
112            let _ = fs::remove_file(&temporary);
113            return Err(ConfigError::Io(error));
114        }
115    }
116    if let Err(error) = temporary_file
117        .write_all(candidate.as_bytes())
118        .and_then(|()| temporary_file.sync_all())
119    {
120        let _ = fs::remove_file(&temporary);
121        return Err(ConfigError::Io(error));
122    }
123    drop(temporary_file);
124    observe(CommitStep::TemporarySynced);
125
126    let unchanged = match fs::read_to_string(path) {
127        Ok(current) => originally_existed && current == original,
128        Err(error) if error.kind() == std::io::ErrorKind::NotFound => !originally_existed,
129        Err(error) => {
130            let _ = fs::remove_file(&temporary);
131            return Err(ConfigError::Io(error));
132        }
133    };
134    if !unchanged {
135        let _ = fs::remove_file(&temporary);
136        return Err(ConfigError::ChangedOnDisk(path.to_owned()));
137    }
138    observe(CommitStep::IdentityChecked);
139    if let Err(error) = replace_file(&temporary, path) {
140        let _ = fs::remove_file(&temporary);
141        return Err(ConfigError::Io(error));
142    }
143    observe(CommitStep::Renamed);
144    sync_directory(parent).map_err(ConfigError::Io)?;
145    observe(CommitStep::ParentSynced);
146    Ok(())
147}
148
149#[cfg(windows)]
150fn create_windows_temporary(
151    parent: &Path,
152    destination: &Path,
153) -> Result<(std::path::PathBuf, fs::File), ConfigError> {
154    use std::{
155        sync::atomic::{AtomicU64, Ordering},
156        time::{SystemTime, UNIX_EPOCH},
157    };
158
159    static SEQUENCE: AtomicU64 = AtomicU64::new(1);
160    let base = destination
161        .file_name()
162        .and_then(|name| name.to_str())
163        .unwrap_or("config.toml");
164    let epoch = SystemTime::now()
165        .duration_since(UNIX_EPOCH)
166        .unwrap_or_default()
167        .as_nanos();
168    for _ in 0..128 {
169        let sequence = SEQUENCE.fetch_add(1, Ordering::Relaxed);
170        let temporary = parent.join(format!(
171            ".{base}.{}.{epoch:x}.{sequence:x}.tmp",
172            std::process::id()
173        ));
174        match cookie_agent_models::secure_store::create_windows_private_file(&temporary) {
175            Ok(file) => return Ok((temporary, file)),
176            Err(error) if error.kind() == std::io::ErrorKind::AlreadyExists => {
177                // CREATE_NEW never follows or truncates a stale file/reparse
178                // point; a collision advances to a fresh unique name.
179            }
180            Err(error) => return Err(ConfigError::Io(error)),
181        }
182    }
183    Err(ConfigError::Io(std::io::Error::new(
184        std::io::ErrorKind::AlreadyExists,
185        "could not allocate a unique private configuration temporary file",
186    )))
187}
188
189#[cfg(unix)]
190fn replace_file(source: &Path, destination: &Path) -> std::io::Result<()> {
191    fs::rename(source, destination)
192}
193
194#[cfg(windows)]
195fn replace_file(source: &Path, destination: &Path) -> std::io::Result<()> {
196    // One shared primitive for every durable Windows replacement: a superseding
197    // POSIX rename, so a concurrent reader opening the destination by name never
198    // finds it absent mid-replacement, with the classic write-through
199    // `MoveFileExW` as the fallback where that rename is unavailable.
200    cookie_agent_models::secure_store::replace_windows_path(source, destination)
201}
202
203#[cfg(unix)]
204fn sync_directory(path: &Path) -> std::io::Result<()> {
205    fs::File::open(path)?.sync_all()
206}
207
208#[cfg(windows)]
209fn sync_directory(path: &Path) -> std::io::Result<()> {
210    use std::os::windows::fs::OpenOptionsExt as _;
211
212    const FILE_FLAG_BACKUP_SEMANTICS: u32 = 0x0200_0000;
213    OpenOptions::new()
214        .read(true)
215        .custom_flags(FILE_FLAG_BACKUP_SEMANTICS)
216        .open(path)
217        .and_then(|directory| match directory.sync_all() {
218            Ok(()) => Ok(()),
219            Err(error) if matches!(error.raw_os_error(), Some(1 | 5 | 6 | 50)) => {
220                // `replace_file` already committed the replacement, and Windows
221                // does not consistently support flushing a directory handle
222                // across filesystems and host policies. This second flush keeps
223                // the `ParentSynced` commit step observable on hosts that do.
224                Ok(())
225            }
226            Err(error) => Err(error),
227        })
228}
229
230fn tempfile_root(path: &Path, candidate: &str) -> Result<std::path::PathBuf, ConfigError> {
231    let parent = path.parent().ok_or(ConfigError::UnsafePath)?;
232    fs::create_dir_all(parent).map_err(ConfigError::Io)?;
233    let root = parent.join(format!(".cookie-agent-validate-{}", std::process::id()));
234    if root.exists() {
235        fs::remove_dir_all(&root).map_err(ConfigError::Io)?;
236    }
237    fs::create_dir(&root).map_err(ConfigError::Io)?;
238    fs::write(root.join("config.toml"), candidate).map_err(ConfigError::Io)?;
239    Ok(root)
240}
241
242#[cfg(test)]
243mod tests {
244    use std::{cell::RefCell, fs};
245
246    use super::{CommitStep, write_mcp_server_observed};
247    use crate::{ConfigError, McpServerConfig};
248
249    fn server(command: &str) -> McpServerConfig {
250        McpServerConfig {
251            command: Some(command.into()),
252            args: Vec::new(),
253            env: Default::default(),
254            cwd: None,
255            url: None,
256            headers: Default::default(),
257            oauth: Default::default(),
258            enabled: true,
259            lazy: false,
260            timeout_ms: None,
261        }
262    }
263
264    #[test]
265    fn durable_replace_orders_sync_identity_rename_and_directory_sync() {
266        let directory = tempfile::tempdir().expect("temporary directory");
267        let path = directory.path().join("config.toml");
268        fs::write(&path, "# retained\n").expect("initial config");
269        let steps = RefCell::new(Vec::new());
270
271        let result = write_mcp_server_observed(&path, "demo", &server("new"), |step| {
272            steps.borrow_mut().push(step);
273            #[cfg(windows)]
274            if step == CommitStep::TemporarySynced {
275                let temporary = fs::read_dir(directory.path())
276                    .expect("list temporary directory")
277                    .filter_map(Result::ok)
278                    .map(|entry| entry.path())
279                    .find(|path| {
280                        path.file_name()
281                            .is_some_and(|name| name.to_string_lossy().ends_with(".tmp"))
282                    })
283                    .expect("unique temporary config");
284                cookie_agent_models::secure_store::verify_windows_private_creation(&temporary)
285                    .expect("temporary config ACL");
286            }
287        });
288        if let Err(error) = result {
289            panic!(
290                "durable write failed: {error}; source: {:?}",
291                std::error::Error::source(&error)
292            );
293        }
294
295        assert_eq!(
296            *steps.borrow(),
297            [
298                CommitStep::TemporarySynced,
299                CommitStep::IdentityChecked,
300                CommitStep::Renamed,
301                CommitStep::ParentSynced,
302            ]
303        );
304        let persisted = fs::read_to_string(&path).expect("read replaced configuration");
305        assert!(
306            persisted.contains("[mcp.servers.demo]"),
307            "replacement content missing from {path:?}: {persisted}"
308        );
309    }
310
311    #[test]
312    fn external_edit_before_commit_is_rejected_without_overwrite() {
313        let directory = tempfile::tempdir().expect("temporary directory");
314        let path = directory.path().join("config.toml");
315        fs::write(&path, "# original\n").expect("initial config");
316
317        let error = write_mcp_server_observed(&path, "demo", &server("new"), |step| {
318            if step == CommitStep::TemporarySynced {
319                fs::write(&path, "# external edit\n").expect("external edit");
320            }
321        })
322        .expect_err("concurrent edit must conflict");
323
324        assert!(matches!(error, ConfigError::ChangedOnDisk(changed) if changed == path));
325        assert_eq!(
326            fs::read_to_string(&path).expect("read preserved external edit"),
327            "# external edit\n"
328        );
329    }
330}
cookie_agent_config/mcp_store.rs

cookie_agent_config/
mcp_store.rs
