use crate::mcp::{
ConnectedServer, JsonField, McpServer, McpTransportType, is_auth_required_error, oauth,
spawn_mcp_server,
};
use anyhow::Result;
use parking_lot::Mutex;
use std::collections::HashMap;
use std::path::Path;
use std::sync::{Arc, Weak};
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub struct McpServerKey {
pub name: String,
pub transport: McpTransportKey,
}
#[derive(Clone, Debug, Eq, Hash, PartialEq)]
pub enum McpTransportKey {
Stdio {
command: String,
args: Vec<String>,
env: Vec<(String, String)>,
},
Remote {
transport_type: McpTransportType,
url: String,
headers: Vec<(String, String)>,
},
}
impl McpServerKey {
pub fn from_spec(name: &str, spec: &McpServer) -> Self {
let transport = if spec.is_remote() {
let url = spec.url.clone().unwrap_or_default();
let mut headers: Vec<(String, String)> = spec
.headers
.as_ref()
.map(|h| h.iter().map(|(k, v)| (k.clone(), v.clone())).collect())
.unwrap_or_default();
headers.sort();
McpTransportKey::Remote {
transport_type: spec.transport_type.clone(),
url,
headers,
}
} else {
let command = spec.command.clone().unwrap_or_default();
let mut args = spec.args.clone().unwrap_or_default();
args.sort();
let mut env: Vec<(String, String)> = spec
.env
.as_ref()
.map(|e| {
e.iter()
.map(|(k, v)| {
let v_str = match v {
JsonField::Str(s) => s.clone(),
JsonField::Bool(b) => b.to_string(),
JsonField::Int(i) => i.to_string(),
};
(k.clone(), v_str)
})
.collect()
})
.unwrap_or_default();
env.sort();
McpTransportKey::Stdio { command, args, env }
};
Self {
name: name.into(),
transport,
}
}
}
#[derive(Default)]
pub struct McpFactory {
active: Mutex<HashMap<McpServerKey, Weak<ConnectedServer>>>,
}
impl McpFactory {
pub fn try_get_active(&self, key: &McpServerKey) -> Option<Arc<ConnectedServer>> {
let map = self.active.lock();
map.get(key).and_then(|weak| weak.upgrade())
}
pub fn insert_active(&self, key: McpServerKey, handle: &Arc<ConnectedServer>) {
let mut map = self.active.lock();
map.insert(key, Arc::downgrade(handle));
}
pub async fn acquire(
&self,
name: &str,
spec: &McpServer,
log_path: Option<&Path>,
) -> Result<Arc<ConnectedServer>> {
let key = McpServerKey::from_spec(name, spec);
if let Some(existing) = self.try_get_active(&key) {
return Ok(existing);
}
let bearer_token = if spec.is_remote() {
oauth::load_valid_mcp_token(name)
} else {
None
};
let handle = spawn_mcp_server(spec, log_path, bearer_token)
.await
.map_err(|e| {
if is_auth_required_error(&e) {
e.context(format!(
"MCP server '{name}' requires OAuth authentication. \
Run `coyote --auth-mcp {name}` or `.mcp auth {name}` in the REPL to authenticate."
))
} else {
e
}
})?;
self.insert_active(key, &handle);
Ok(handle)
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::mcp::{JsonField, McpServer, McpTransportType};
use indexmap::IndexMap;
use std::collections::HashMap;
fn stdio_spec(
command: &str,
args: Option<Vec<String>>,
env: Option<IndexMap<String, JsonField>>,
) -> McpServer {
McpServer {
transport_type: McpTransportType::Stdio,
command: Some(command.to_string()),
args,
env,
cwd: None,
url: None,
headers: None,
oauth: None,
}
}
fn remote_spec(
transport: McpTransportType,
url: &str,
headers: Option<IndexMap<String, String>>,
) -> McpServer {
McpServer {
transport_type: transport,
command: None,
args: None,
env: None,
cwd: None,
url: Some(url.to_string()),
headers,
oauth: None,
}
}
#[test]
fn key_from_stdio_spec_captures_command_args_env() {
let mut env = IndexMap::new();
env.insert("TOKEN".into(), JsonField::Str("abc".into()));
let spec = stdio_spec("npx", Some(vec!["-y".into(), "server".into()]), Some(env));
let key = McpServerKey::from_spec("my-server", &spec);
assert_eq!(key.name, "my-server");
match &key.transport {
McpTransportKey::Stdio { command, args, env } => {
assert_eq!(command, "npx");
assert_eq!(args, &["-y", "server"]);
assert_eq!(env, &[("TOKEN".to_string(), "abc".to_string())]);
}
_ => panic!("expected Stdio transport key"),
}
}
#[test]
fn key_from_stdio_spec_sorts_args_and_env() {
let mut env = IndexMap::new();
env.insert("Z_VAR".into(), JsonField::Str("z".into()));
env.insert("A_VAR".into(), JsonField::Int(42));
let spec = stdio_spec(
"cmd",
Some(vec!["charlie".into(), "alpha".into(), "bravo".into()]),
Some(env),
);
let key = McpServerKey::from_spec("s", &spec);
match &key.transport {
McpTransportKey::Stdio { args, env, .. } => {
assert_eq!(args, &["alpha", "bravo", "charlie"]);
assert_eq!(env[0].0, "A_VAR");
assert_eq!(env[0].1, "42");
assert_eq!(env[1].0, "Z_VAR");
assert_eq!(env[1].1, "z");
}
_ => panic!("expected Stdio"),
}
}
#[test]
fn key_from_stdio_spec_defaults_empty_when_none() {
let spec = stdio_spec("echo", None, None);
let key = McpServerKey::from_spec("bare", &spec);
match &key.transport {
McpTransportKey::Stdio { command, args, env } => {
assert_eq!(command, "echo");
assert!(args.is_empty());
assert!(env.is_empty());
}
_ => panic!("expected Stdio"),
}
}
#[test]
fn key_from_remote_http_spec() {
let spec = remote_spec(McpTransportType::Http, "http://localhost:8080", None);
let key = McpServerKey::from_spec("http-srv", &spec);
assert_eq!(key.name, "http-srv");
match &key.transport {
McpTransportKey::Remote {
transport_type,
url,
headers,
} => {
assert_eq!(*transport_type, McpTransportType::Http);
assert_eq!(url, "http://localhost:8080");
assert!(headers.is_empty());
}
_ => panic!("expected Remote"),
}
}
#[test]
fn key_from_remote_sse_spec_with_sorted_headers() {
let mut hdrs = IndexMap::new();
hdrs.insert("Z-Key".into(), "z-val".into());
hdrs.insert("A-Key".into(), "a-val".into());
let spec = remote_spec(McpTransportType::Sse, "http://sse.example.com", Some(hdrs));
let key = McpServerKey::from_spec("sse-srv", &spec);
match &key.transport {
McpTransportKey::Remote { headers, .. } => {
assert_eq!(headers[0], ("A-Key".to_string(), "a-val".to_string()));
assert_eq!(headers[1], ("Z-Key".to_string(), "z-val".to_string()));
}
_ => panic!("expected Remote"),
}
}
#[test]
fn key_equality_same_spec_produces_equal_keys() {
let spec = stdio_spec("npx", Some(vec!["a".into()]), None);
let k1 = McpServerKey::from_spec("s", &spec);
let k2 = McpServerKey::from_spec("s", &spec);
assert_eq!(k1, k2);
}
#[test]
fn key_inequality_different_names() {
let spec = stdio_spec("npx", None, None);
let k1 = McpServerKey::from_spec("a", &spec);
let k2 = McpServerKey::from_spec("b", &spec);
assert_ne!(k1, k2);
}
#[test]
fn key_inequality_different_commands() {
let s1 = stdio_spec("npx", None, None);
let s2 = stdio_spec("node", None, None);
let k1 = McpServerKey::from_spec("s", &s1);
let k2 = McpServerKey::from_spec("s", &s2);
assert_ne!(k1, k2);
}
#[test]
fn key_env_bool_and_int_coerce_to_string() {
let mut env = IndexMap::new();
env.insert("FLAG".into(), JsonField::Bool(true));
env.insert("PORT".into(), JsonField::Int(3000));
let spec = stdio_spec("cmd", None, Some(env));
let key = McpServerKey::from_spec("s", &spec);
match &key.transport {
McpTransportKey::Stdio { env, .. } => {
let map: HashMap<&str, &str> =
env.iter().map(|(k, v)| (k.as_str(), v.as_str())).collect();
assert_eq!(map["FLAG"], "true");
assert_eq!(map["PORT"], "3000");
}
_ => panic!("expected Stdio"),
}
}
#[test]
fn factory_try_get_active_returns_none_when_empty() {
let factory = McpFactory::default();
let spec = stdio_spec("cmd", None, None);
let key = McpServerKey::from_spec("s", &spec);
assert!(factory.try_get_active(&key).is_none());
}
#[test]
fn factory_try_get_active_returns_none_for_unknown_key() {
let factory = McpFactory::default();
let spec = stdio_spec("cmd", None, None);
let key = McpServerKey::from_spec("s", &spec);
assert!(factory.try_get_active(&key).is_none());
}
#[test]
fn factory_default_has_empty_active_map() {
let factory = McpFactory::default();
let map = factory.active.lock();
assert!(map.is_empty());
}
}