aion-worker-mcp 0.31.0

//! The `mcpServers` document — the same JSON Claude Code, Cursor and the MCP
//! inspector read, so an operator's existing configuration drops in as is.

use std::collections::BTreeMap;

use serde::Deserialize;

/// The whole document: a map of server name to how to reach it.
#[derive(Debug, Deserialize)]
pub struct McpConfig {
    /// Servers by name. The name becomes the activity prefix `<name>.<tool>`.
    #[serde(rename = "mcpServers")]
    pub mcp_servers: BTreeMap<String, ServerSpec>,
}

/// One server entry. Unknown keys (`type`, `disabled`, editor extras) are
/// tolerated because the document is shared with other clients; the fields
/// this worker acts on are these.
#[derive(Debug, Clone, Deserialize)]
pub struct ServerSpec {
    /// Executable for a stdio server.
    pub command: Option<String>,
    /// Arguments for the stdio server.
    #[serde(default)]
    pub args: Vec<String>,
    /// Extra environment for the stdio server (added to the worker's own).
    #[serde(default)]
    pub env: BTreeMap<String, String>,
    /// Working directory for the stdio server.
    pub cwd: Option<String>,
    /// URL of a remote (Streamable HTTP) server. Not served by this worker yet;
    /// named as a refusal at startup, never skipped in silence.
    pub url: Option<String>,
    /// Some documents mark a server off without deleting it.
    #[serde(default)]
    pub disabled: bool,
}

/// Why a document cannot be served.
#[derive(Debug, thiserror::Error)]
pub enum ConfigError {
    /// The text is not the document.
    #[error("not an mcpServers document: {0}")]
    Json(#[from] serde_json::Error),
    /// Nothing to serve.
    #[error("the document names no enabled servers")]
    Empty,
    /// A remote server: this worker speaks stdio only, and says so.
    #[error("server `{name}` is a remote server ({url}); this worker serves stdio servers only")]
    RemoteNotServed {
        /// Server name in the document.
        name: String,
        /// Its URL.
        url: String,
    },
    /// Neither a command nor a URL.
    #[error("server `{name}` has neither `command` nor `url`")]
    NoTransport {
        /// Server name in the document.
        name: String,
    },
    /// A name this worker cannot spell as an activity prefix.
    #[error(
        "server name `{name}` may not contain `.`, whitespace or `/`: it prefixes activity names"
    )]
    BadName {
        /// The offending name.
        name: String,
    },
}

impl McpConfig {
    /// Parses and validates the document: every enabled server is a stdio
    /// server with a command and a name that can prefix an activity.
    ///
    /// # Errors
    ///
    /// Returns the first reason the document cannot be served, named.
    pub fn parse(text: &str) -> Result<Self, ConfigError> {
        let parsed: Self = serde_json::from_str(text)?;
        let enabled = parsed.enabled_servers();
        if enabled.is_empty() {
            return Err(ConfigError::Empty);
        }
        for (name, spec) in enabled {
            if !is_identifier(name) {
                return Err(ConfigError::BadName { name: name.clone() });
            }
            match (&spec.command, &spec.url) {
                (Some(_), _) => {}
                (None, Some(url)) => {
                    return Err(ConfigError::RemoteNotServed {
                        name: name.clone(),
                        url: url.clone(),
                    });
                }
                (None, None) => return Err(ConfigError::NoTransport { name: name.clone() }),
            }
        }
        Ok(parsed)
    }

    /// Servers that are on.
    pub fn enabled_servers(&self) -> Vec<(&String, &ServerSpec)> {
        self.mcp_servers
            .iter()
            .filter(|(_, spec)| !spec.disabled)
            .collect()
    }
}

/// An AWL identifier: what `worker <name>` and `action <name>` accept.
pub fn is_identifier(text: &str) -> bool {
    let mut chars = text.chars();
    match chars.next() {
        Some(first) if first.is_ascii_alphabetic() || first == '_' => {
            chars.all(|c| c.is_ascii_alphanumeric() || c == '_')
        }
        _ => false,
    }
}

/// The identifier a tool is served as: every character an identifier cannot
/// carry becomes `_`, and a leading digit is prefixed. `get-sum` is served as
/// `get_sum`; the registration log names both.
pub fn activity_name(tool: &str) -> String {
    let mut out: String = tool
        .chars()
        .map(|c| {
            if c.is_ascii_alphanumeric() || c == '_' {
                c
            } else {
                '_'
            }
        })
        .collect();
    if out.chars().next().is_some_and(|c| c.is_ascii_digit()) {
        out.insert(0, '_');
    }
    if out.is_empty() { "_".to_owned() } else { out }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn a_claude_code_document_parses_with_its_extra_keys() -> Result<(), ConfigError> {
        let doc = r#"{"mcpServers":{"fs":{"command":"npx","args":["-y","x"],"env":{"A":"1"},"type":"stdio"},
                      "off":{"command":"x","disabled":true}}}"#;
        let config = McpConfig::parse(doc)?;
        let enabled = config.enabled_servers();
        assert_eq!(enabled.len(), 1);
        assert_eq!(enabled[0].0, "fs");
        assert_eq!(enabled[0].1.args, vec!["-y", "x"]);
        assert_eq!(enabled[0].1.env.get("A").map(String::as_str), Some("1"));
        Ok(())
    }

    #[test]
    fn a_remote_server_is_refused_by_name() {
        let doc = r#"{"mcpServers":{"remote":{"url":"https://example.test/mcp"}}}"#;
        let outcome = McpConfig::parse(doc);
        assert!(
            matches!(outcome, Err(ConfigError::RemoteNotServed { ref name, .. }) if name == "remote"),
            "expected a named refusal, got {outcome:?}"
        );
    }

    #[test]
    fn an_empty_or_all_disabled_document_is_refused() {
        assert!(matches!(
            McpConfig::parse(r#"{"mcpServers":{}}"#),
            Err(ConfigError::Empty)
        ));
        let doc = r#"{"mcpServers":{"a":{"command":"x","disabled":true}}}"#;
        assert!(matches!(McpConfig::parse(doc), Err(ConfigError::Empty)));
    }

    #[test]
    fn a_name_that_is_not_an_identifier_is_refused() {
        for bad in ["a.b", "a-b", "1a", "", "a b"] {
            let doc = format!(r#"{{"mcpServers":{{"{bad}":{{"command":"x"}}}}}}"#);
            assert!(
                matches!(McpConfig::parse(&doc), Err(ConfigError::BadName { .. })),
                "{bad}"
            );
        }
        assert!(McpConfig::parse(r#"{"mcpServers":{"fs_1":{"command":"x"}}}"#).is_ok());
    }

    #[test]
    fn tool_names_become_identifiers_predictably() {
        assert_eq!(activity_name("get-sum"), "get_sum");
        assert_eq!(activity_name("get.sum/2"), "get_sum_2");
        assert_eq!(activity_name("1st"), "_1st");
        assert_eq!(activity_name("plain_name"), "plain_name");
    }

    #[test]
    fn a_server_with_no_transport_is_refused() {
        let doc = r#"{"mcpServers":{"a":{"args":["x"]}}}"#;
        assert!(matches!(
            McpConfig::parse(doc),
            Err(ConfigError::NoTransport { .. })
        ));
    }
}