1use std::{collections::BTreeMap, path::PathBuf};
2
3use cookie_agent_config::{
4 LoadedConfiguration, LoadedMcpServer, McpOAuthConfig, McpOAuthSettings, McpServerConfig,
5 McpServerSource, write_mcp_server,
6};
7use cookie_agent_protocol::{
8 McpConfigSource, McpConfigTarget, McpOAuthDefinition, McpOAuthSettingsDefinition,
9 McpServerDefinition, McpServerInfo, McpServerState as WireMcpServerState,
10};
11
12use crate::{Engine, EngineError, McpServerState};
13
14#[derive(Clone, Debug)]
15pub(crate) struct ConfigStore {
16 user: BTreeMap<String, McpServerConfig>,
17 workspace: BTreeMap<String, McpServerConfig>,
18 runtime: BTreeMap<String, Option<McpServerConfig>>,
19 user_path: Option<PathBuf>,
20 workspace_path: Option<PathBuf>,
21}
22
23impl ConfigStore {
24 pub(crate) fn new(config: &LoadedConfiguration) -> Self {
25 Self {
26 user: config.user_mcp_servers.clone(),
27 workspace: config.workspace_mcp_servers.clone(),
28 runtime: BTreeMap::new(),
29 user_path: config.config_paths.user.clone(),
30 workspace_path: config.config_paths.workspace.clone(),
31 }
32 }
33
34 fn effective(&self, name: &str) -> Option<LoadedMcpServer> {
35 if let Some(runtime) = self.runtime.get(name) {
36 return runtime.clone().map(|config| LoadedMcpServer {
37 source: McpServerSource::Runtime,
38 config,
39 });
40 }
41 self.workspace
42 .get(name)
43 .cloned()
44 .map(|config| LoadedMcpServer {
45 source: McpServerSource::WorkspaceFile,
46 config,
47 })
48 .or_else(|| {
49 self.user.get(name).cloned().map(|config| LoadedMcpServer {
50 source: McpServerSource::UserFile,
51 config,
52 })
53 })
54 }
55
56 fn effective_all(&self) -> BTreeMap<String, LoadedMcpServer> {
57 let names = self
58 .user
59 .keys()
60 .chain(self.workspace.keys())
61 .chain(self.runtime.keys())
62 .cloned()
63 .collect::<std::collections::BTreeSet<_>>();
64 names
65 .into_iter()
66 .filter_map(|name| self.effective(&name).map(|server| (name, server)))
67 .collect()
68 }
69
70 fn add(&mut self, name: &str, config: McpServerConfig) -> Result<(), EngineError> {
71 if self.effective(name).is_some() {
72 return Err(EngineError::Mcp(format!("server `{name}` already exists")));
73 }
74 config
75 .validate(name)
76 .map_err(|error| EngineError::Mcp(error.to_string()))?;
77 self.runtime.insert(name.to_owned(), Some(config));
78 Ok(())
79 }
80
81 fn edit(&mut self, name: &str, config: McpServerConfig) -> Result<(), EngineError> {
82 if self.effective(name).is_none() {
83 return Err(EngineError::Mcp(format!("unknown server `{name}`")));
84 }
85 config
86 .validate(name)
87 .map_err(|error| EngineError::Mcp(error.to_string()))?;
88 self.runtime.insert(name.to_owned(), Some(config));
89 Ok(())
90 }
91
92 fn remove(&mut self, name: &str) -> Result<(), EngineError> {
93 if self.effective(name).is_none() {
94 return Err(EngineError::Mcp(format!("unknown server `{name}`")));
95 }
96 self.runtime.insert(name.to_owned(), None);
97 Ok(())
98 }
99
100 fn set_enabled(&mut self, name: &str, enabled: bool) -> Result<(), EngineError> {
101 let mut server = self
102 .effective(name)
103 .ok_or_else(|| EngineError::Mcp(format!("unknown server `{name}`")))?
104 .config;
105 server.enabled = enabled;
106 self.runtime.insert(name.to_owned(), Some(server));
107 Ok(())
108 }
109
110 fn persist(&mut self, name: &str, target: McpConfigTarget) -> Result<(), EngineError> {
111 let config = self
112 .runtime
113 .get(name)
114 .and_then(Option::as_ref)
115 .cloned()
116 .ok_or_else(|| {
117 EngineError::Mcp(format!("server `{name}` has no runtime entry to persist"))
118 })?;
119 let (path, layer) = match target {
120 McpConfigTarget::UserFile => (&self.user_path, &mut self.user),
121 McpConfigTarget::WorkspaceFile => (&self.workspace_path, &mut self.workspace),
122 };
123 let path = path.as_ref().ok_or_else(|| {
124 EngineError::Mcp("the selected configuration file layer is unavailable".into())
125 })?;
126 write_mcp_server(path, name, &config)
127 .map_err(|error| EngineError::Mcp(error.to_string()))?;
128 layer.insert(name.to_owned(), config);
129 Ok(())
130 }
131}
132
133impl Engine {
134 #[must_use]
135 pub fn list_mcp_servers(&self) -> Vec<McpServerInfo> {
136 let store = self
137 .inner
138 .config_store
139 .lock()
140 .unwrap_or_else(|poisoned| poisoned.into_inner());
141 let statuses = self
142 .inner
143 .mcp
144 .statuses()
145 .into_iter()
146 .map(|status| (status.server.clone(), status))
147 .collect::<BTreeMap<_, _>>();
148 store
149 .effective_all()
150 .into_iter()
151 .map(|(name, loaded)| {
152 let status = statuses.get(&name);
153 let state = status.map_or(WireMcpServerState::Disconnected, |status| match status
154 .state
155 {
156 McpServerState::Disabled => WireMcpServerState::Disabled,
157 McpServerState::Disconnected if loaded.config.lazy => {
158 WireMcpServerState::LazyNotConnected
159 }
160 McpServerState::Disconnected => WireMcpServerState::Disconnected,
161 McpServerState::Connecting => WireMcpServerState::Connecting,
162 McpServerState::Connected => WireMcpServerState::Connected,
163 McpServerState::NeedsAuth => WireMcpServerState::NeedsAuth,
164 McpServerState::Failed => WireMcpServerState::Failed,
165 });
166 McpServerInfo {
167 name,
168 source: source_to_wire(loaded.source),
169 definition: definition_to_wire(&loaded.config),
170 state,
171 tool_count: status.map_or(0, |status| status.tools.len() as u32),
172 message: status.and_then(|status| status.message.clone()),
173 auth_in_progress: status.map(|status| status.auth_in_progress),
174 }
175 })
176 .collect()
177 }
178
179 pub async fn add_mcp_server(
180 &self,
181 name: String,
182 definition: McpServerDefinition,
183 ) -> Result<Option<McpServerInfo>, EngineError> {
184 self.mutate_mcp(&name, |store| {
185 store.add(&name, definition_from_wire(definition))
186 })
187 .await
188 }
189
190 pub async fn edit_mcp_server(
191 &self,
192 name: String,
193 definition: McpServerDefinition,
194 ) -> Result<Option<McpServerInfo>, EngineError> {
195 self.mutate_mcp(&name, |store| {
196 store.edit(&name, definition_from_wire(definition))
197 })
198 .await
199 }
200
201 pub async fn remove_mcp_server(
202 &self,
203 name: String,
204 ) -> Result<Option<McpServerInfo>, EngineError> {
205 self.mutate_mcp(&name, |store| store.remove(&name)).await
206 }
207
208 pub async fn set_mcp_server_enabled(
209 &self,
210 name: String,
211 enabled: bool,
212 ) -> Result<Option<McpServerInfo>, EngineError> {
213 self.mutate_mcp(&name, |store| store.set_enabled(&name, enabled))
214 .await
215 }
216
217 pub async fn reconnect_mcp_server(
218 &self,
219 name: String,
220 ) -> Result<Option<McpServerInfo>, EngineError> {
221 let _mutation = self.inner.mcp_mutation.lock().await;
222 self.inner
223 .mcp
224 .reconnect_server(&name)
225 .await
226 .map_err(|error| EngineError::Mcp(error.to_string()))?;
227 Ok(self
228 .list_mcp_servers()
229 .into_iter()
230 .find(|server| server.name == name))
231 }
232
233 pub async fn begin_mcp_auth(&self, name: String) -> Result<String, EngineError> {
234 let _mutation = self.inner.mcp_mutation.lock().await;
235 self.inner
236 .mcp
237 .begin_auth(&name)
238 .await
239 .map_err(|error| EngineError::Mcp(error.to_string()))
240 }
241
242 pub async fn cancel_mcp_auth(&self, name: String) -> Result<(), EngineError> {
243 let _mutation = self.inner.mcp_mutation.lock().await;
244 self.inner
245 .mcp
246 .cancel_auth(&name)
247 .map_err(|error| EngineError::Mcp(error.to_string()))
248 }
249
250 pub async fn persist_mcp_server(
251 &self,
252 name: String,
253 target: McpConfigTarget,
254 ) -> Result<Option<McpServerInfo>, EngineError> {
255 let _mutation = self.inner.mcp_mutation.lock().await;
256 self.inner
257 .config_store
258 .lock()
259 .unwrap_or_else(|poisoned| poisoned.into_inner())
260 .persist(&name, target)?;
261 Ok(self
262 .list_mcp_servers()
263 .into_iter()
264 .find(|server| server.name == name))
265 }
266
267 async fn mutate_mcp(
268 &self,
269 name: &str,
270 mutation: impl FnOnce(&mut ConfigStore) -> Result<(), EngineError>,
271 ) -> Result<Option<McpServerInfo>, EngineError> {
272 let _mutation = self.inner.mcp_mutation.lock().await;
273 let previous = {
274 let mut store = self
275 .inner
276 .config_store
277 .lock()
278 .unwrap_or_else(|poisoned| poisoned.into_inner());
279 let previous = store.clone();
280 mutation(&mut store)?;
281 previous
282 };
283 let effective = self
284 .inner
285 .config_store
286 .lock()
287 .unwrap_or_else(|poisoned| poisoned.into_inner())
288 .effective(name);
289 let result = match effective {
290 Some(server) => self.inner.mcp.upsert_server(name.to_owned(), server).await,
291 None => self.inner.mcp.remove_server(name).await,
292 };
293 if let Err(error) = result {
294 *self
295 .inner
296 .config_store
297 .lock()
298 .unwrap_or_else(|poisoned| poisoned.into_inner()) = previous;
299 return Err(EngineError::Mcp(error.to_string()));
300 }
301 Ok(self
302 .list_mcp_servers()
303 .into_iter()
304 .find(|server| server.name == name))
305 }
306}
307
308fn source_to_wire(source: McpServerSource) -> McpConfigSource {
309 match source {
310 McpServerSource::UserFile => McpConfigSource::UserFile,
311 McpServerSource::WorkspaceFile => McpConfigSource::WorkspaceFile,
312 McpServerSource::Runtime => McpConfigSource::Runtime,
313 }
314}
315
316fn definition_to_wire(config: &McpServerConfig) -> McpServerDefinition {
317 McpServerDefinition {
318 command: config.command.clone(),
319 args: config.args.clone(),
320 env: config.env.clone(),
321 cwd: config.cwd.clone(),
322 url: config.url.clone(),
323 headers: config.headers.clone(),
324 oauth: Some(match &config.oauth {
325 McpOAuthConfig::Disabled => McpOAuthDefinition::Bool(false),
326 McpOAuthConfig::Auto | McpOAuthConfig::Enabled => McpOAuthDefinition::Bool(true),
327 McpOAuthConfig::Settings(settings) => {
328 McpOAuthDefinition::Settings(McpOAuthSettingsDefinition {
329 client_id: settings.client_id.clone(),
330 client_secret: settings.client_secret.clone(),
331 client_metadata_url: settings.client_metadata_url.clone(),
332 scopes: (!settings.scopes.is_empty()).then(|| settings.scopes.clone()),
333 })
334 }
335 }),
336 enabled: config.enabled,
337 lazy: config.lazy,
338 timeout_ms: config.timeout_ms,
339 }
340}
341
342fn definition_from_wire(definition: McpServerDefinition) -> McpServerConfig {
343 let stdio = definition.command.is_some();
344 McpServerConfig {
345 command: definition.command,
346 args: definition.args,
347 env: definition.env,
348 cwd: definition.cwd,
349 url: definition.url,
350 headers: definition.headers,
351 oauth: match definition.oauth {
352 None => McpOAuthConfig::Auto,
353 Some(McpOAuthDefinition::Bool(true)) if stdio => McpOAuthConfig::Auto,
354 Some(McpOAuthDefinition::Bool(true)) => McpOAuthConfig::Enabled,
355 Some(McpOAuthDefinition::Bool(false)) => McpOAuthConfig::Disabled,
356 Some(McpOAuthDefinition::Settings(settings)) => {
357 McpOAuthConfig::Settings(McpOAuthSettings {
358 client_id: settings.client_id,
359 client_secret: settings.client_secret,
360 client_metadata_url: settings.client_metadata_url,
361 scopes: settings.scopes.unwrap_or_default(),
362 })
363 }
364 },
365 enabled: definition.enabled,
366 lazy: definition.lazy,
367 timeout_ms: definition.timeout_ms,
368 }
369}
370
371#[cfg(test)]
372mod tests {
373 use std::fs;
374
375 use cookie_agent_config::{McpServerSource, load_from_roots};
376 use cookie_agent_protocol::McpConfigTarget;
377
378 use super::ConfigStore;
379
380 #[test]
381 fn runtime_layer_overrides_files_can_remove_and_writes_back_strictly() {
382 let user = tempfile::tempdir().expect("user root");
383 let workspace = tempfile::tempdir().expect("workspace root");
384 fs::write(
385 user.path().join("config.toml"),
386 "# user comment\n[mcp.servers.demo]\ncommand = \"user\"\n",
387 )
388 .expect("user config");
389 fs::write(
390 workspace.path().join("config.toml"),
391 "[mcp.servers.demo]\ncommand = \"workspace\"\n",
392 )
393 .expect("workspace config");
394 let loaded = load_from_roots(Some(user.path()), Some(workspace.path())).expect("layers");
395 let mut store = ConfigStore::new(&loaded);
396 assert_eq!(
397 store.effective("demo").expect("workspace server").source,
398 McpServerSource::WorkspaceFile
399 );
400 let mut runtime = store.effective("demo").expect("server").config;
401 runtime.command = Some("runtime".into());
402 store.edit("demo", runtime.clone()).expect("runtime edit");
403 assert_eq!(
404 store.effective("demo").expect("runtime server").source,
405 McpServerSource::Runtime
406 );
407 store
408 .persist("demo", McpConfigTarget::UserFile)
409 .expect("write user layer");
410 assert!(
411 fs::read_to_string(user.path().join("config.toml"))
412 .expect("written user config")
413 .contains("runtime")
414 );
415 store.remove("demo").expect("runtime tombstone");
416 assert!(store.effective("demo").is_none());
417 let reloaded = load_from_roots(Some(user.path()), Some(workspace.path()))
418 .expect("strictly reloaded file layers");
419 assert_eq!(
420 reloaded.mcp_servers["demo"].source,
421 McpServerSource::WorkspaceFile
422 );
423 }
424}