tailscale_mcp/

registry.rs

//! The tool table, and the macro that fills it.
//!
//! Every tool is declared once. That declaration produces its parameter type's
//! schema, the shim that deserialises into it, and the metadata row — so there
//! is no way to add a tool and forget to classify it, and no second list to
//! fall out of step with the first.
//!
//! The router is this table. `list_tools` filters it through the
//! [`crate::gating::Gate`]; `call_tool` looks a name up in it. There is
//! no parallel registration anywhere.

use std::collections::BTreeMap;
use std::sync::Arc;

use rmcp::model::{JsonObject, Tool, ToolAnnotations};
use serde_json::Value;
use tailscale_cli::BoxFuture;

use crate::context::ToolContext;
use crate::error::{ToolError, ToolResult};
use crate::gating::Gate;
use crate::meta::ToolMeta;

/// The name of the parameter that carries the caller's intent.
///
/// Injected into the schema of every tool that requires it, rather than being
/// repeated in twenty parameter structs, so that the rule and its enforcement
/// are the same line of code.
pub const CONFIRM_PARAM: &str = "confirm";

/// The longest name the protocol permits.
pub const MAX_NAME_LEN: usize = 128;

/// How a tool is actually run.
pub type InvokeFn = fn(Arc<ToolContext>, JsonObject) -> BoxFuture<'static, ToolResult<Value>>;

/// One tool: its row in the table, its schema, and its handler.
#[derive(Clone)]
pub struct ToolEntry {
    pub meta: ToolMeta,
    /// Built on demand and cached by the SDK, keyed on the parameter type.
    pub schema: fn() -> Result<Arc<JsonObject>, String>,
    pub invoke: InvokeFn,
}

impl std::fmt::Debug for ToolEntry {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ToolEntry")
            .field("meta", &self.meta)
            .finish_non_exhaustive()
    }
}

impl ToolEntry {
    /// The tool as a client sees it.
    ///
    /// Two things are added here rather than at the declaration: the
    /// annotations, which are derived from the tier so they cannot contradict
    /// it, and the confirmation parameter, which is added to exactly the tools
    /// whose row asks for it.
    pub fn describe(&self) -> Result<Tool, String> {
        let mut schema = (self.schema)()?.as_ref().clone();
        if self.meta.requires_confirmation {
            add_confirm_property(&mut schema);
        }

        let a = self.meta.annotations();
        Ok(
            Tool::new(self.meta.name, self.meta.summary, Arc::new(schema)).with_annotations(
                ToolAnnotations::new()
                    .read_only(a.read_only)
                    .destructive(a.destructive)
                    .idempotent(a.idempotent)
                    .open_world(a.open_world),
            ),
        )
    }
}

/// Add the confirmation flag to a generated schema.
fn add_confirm_property(schema: &mut JsonObject) {
    let properties = schema
        .entry("properties")
        .or_insert_with(|| Value::Object(JsonObject::new()));
    if let Some(properties) = properties.as_object_mut() {
        properties.insert(
            CONFIRM_PARAM.to_owned(),
            serde_json::json!({
                "type": "boolean",
                "default": false,
                "description":
                    "Set to true to confirm this operation. It is irreversible, \
                     or it can disconnect the node this server runs on, so it \
                     will not run without an explicit intent.",
            }),
        );
    }
}

/// The tool table.
///
/// Built once from the declarations, then shared. Lookup is by name because
/// that is what a call carries; listing walks the ordered vector so that the
/// output is stable between runs.
#[derive(Debug, Clone)]
pub struct Registry {
    entries: Vec<ToolEntry>,
    by_name: BTreeMap<&'static str, usize>,
}

impl Registry {
    /// Build a registry, rejecting a table that could not be served correctly.
    ///
    /// The checks are cheap and run at startup rather than in a test alone,
    /// because a duplicate name silently shadowing a tool is the kind of bug
    /// that survives a release.
    pub fn new(entries: Vec<ToolEntry>) -> Result<Self, RegistryError> {
        let mut by_name = BTreeMap::new();
        for (index, entry) in entries.iter().enumerate() {
            let name = entry.meta.name;
            validate_name(name)?;
            if by_name.insert(name, index).is_some() {
                return Err(RegistryError::DuplicateName(name));
            }
            if entry.meta.self_severing && !entry.meta.requires_confirmation {
                return Err(RegistryError::SelfSeveringWithoutConfirmation(name));
            }
            match (entry.schema)() {
                Err(reason) => return Err(RegistryError::BadSchema { name, reason }),
                // A tool that severs when it is aimed at us has to give the
                // caller somewhere to say so. The flag alone would be a claim
                // the schema does not keep (Q83).
                Ok(schema) if entry.meta.severs_local_node => {
                    let asks = schema
                        .get("properties")
                        .and_then(Value::as_object)
                        .is_some_and(|properties| properties.contains_key(CONFIRM_PARAM));
                    if !asks {
                        return Err(RegistryError::SeversLocalWithoutConfirmParameter(name));
                    }
                }
                Ok(_) => {}
            }
        }
        Ok(Self { entries, by_name })
    }

    /// The metadata rows, for anything that needs the table without the
    /// handlers: the gate's emptiness check, the `tools` subcommand, the
    /// generated documentation.
    pub fn metas(&self) -> Vec<ToolMeta> {
        self.entries.iter().map(|e| e.meta).collect()
    }

    pub fn get(&self, name: &str) -> Option<&ToolEntry> {
        self.by_name.get(name).map(|index| &self.entries[*index])
    }

    pub fn len(&self) -> usize {
        self.entries.len()
    }

    pub fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    /// The tools this server offers, in name order.
    ///
    /// A tool the gate does not permit is absent, not present-and-failing: a
    /// model that can see a tool will reach for it.
    pub fn visible(&self, gate: &Gate) -> Vec<&ToolEntry> {
        let mut visible: Vec<&ToolEntry> = self
            .entries
            .iter()
            .filter(|entry| gate.permits(&entry.meta))
            .collect();
        visible.sort_by_key(|entry| entry.meta.name);
        visible
    }

    /// Look a call up, applying the gate and the confirmation rule.
    ///
    /// Returns the entry and the arguments with the confirmation flag removed,
    /// so a handler never has to know the flag exists.
    pub fn resolve(
        &self,
        name: &str,
        mut args: JsonObject,
        gate: &Gate,
    ) -> ToolResult<(&ToolEntry, JsonObject)> {
        let entry = self
            .get(name)
            .ok_or_else(|| ToolError::not_found(&format!("the tool `{name}`")))?;

        if !gate.permits(&entry.meta) {
            return Err(ToolError::not_permitted(name, &gate.needs(&entry.meta)));
        }

        if entry.meta.requires_confirmation {
            let confirmed = match args.remove(CONFIRM_PARAM) {
                Some(Value::Bool(confirmed)) => confirmed,
                None | Some(Value::Null) => false,
                Some(other) => {
                    return Err(ToolError::invalid_args(format!(
                        "`{CONFIRM_PARAM}` must be true or false, not {other}"
                    )));
                }
            };
            if !confirmed {
                return Err(ToolError::confirmation_required(
                    name,
                    confirmation_consequence(&entry.meta),
                ));
            }
        }

        Ok((entry, args))
    }
}

/// The clause that follows the tool name in a confirmation message.
fn confirmation_consequence(meta: &crate::meta::ToolMeta) -> &'static str {
    if meta.self_severing {
        "can disconnect the node this server runs on, so it needs `confirm: true`"
    } else {
        "cannot be undone, so it needs `confirm: true`"
    }
}

/// A table that could not be served.
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum RegistryError {
    #[error("two tools are both named `{0}`")]
    DuplicateName(&'static str),

    #[error("the tool name `{name}` {reason}")]
    BadName {
        name: &'static str,
        reason: &'static str,
    },

    #[error("`{0}` is self-severing but does not require confirmation")]
    SelfSeveringWithoutConfirmation(&'static str),

    #[error("`{0}` severs the local node but its parameters carry no `confirm`")]
    SeversLocalWithoutConfirmParameter(&'static str),

    #[error("the parameters of `{name}` do not make a valid input schema: {reason}")]
    BadSchema { name: &'static str, reason: String },
}

/// The protocol's rules for a tool name, plus ours.
fn validate_name(name: &'static str) -> Result<(), RegistryError> {
    let bad = |reason| Err(RegistryError::BadName { name, reason });
    if name.is_empty() {
        return bad("is empty");
    }
    if name.len() > MAX_NAME_LEN {
        return bad("is longer than the 128 characters the protocol allows");
    }
    if !name
        .bytes()
        .all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-' || b == b'.')
    {
        return bad("contains a character outside [A-Za-z0-9_.-]");
    }
    if !name.starts_with("tailscale_") && !name.starts_with("tailnet_") {
        return bad("does not begin with `tailscale_` or `tailnet_`");
    }
    Ok(())
}

/// Declare tools.
///
/// One block per toolset module; one entry per tool. The entry names its
/// parameter type and its handler, and states its toolset and tier — omit
/// either and this does not compile, which is the point.
///
/// ```ignore
/// tools! {
///     /// Report the state of the local node and the peers it knows about.
///     tailscale_status => StatusParams, handlers::status,
///         toolset: LocalStatus, tier: Read, idempotent: true;
/// }
/// ```
///
/// Optional trailing settings, each defaulting to the safe answer:
/// `idempotent` (false), `confirm` (false), `severing` (false, and implies
/// `confirm`), `severs_local` (false, and requires a `confirm` parameter),
/// `varying` (false), `since` (no minimum version).
#[macro_export]
macro_rules! tools {
    (
        $(
            $(#[doc = $summary:literal])+
            $name:ident => $params:ty, $handler:path,
                toolset: $toolset:ident,
                tier: $tier:ident
                $(, idempotent: $idempotent:literal)?
                $(, confirm: $confirm:literal)?
                $(, severing: $severing:literal)?
                $(, severs_local: $severs_local:literal)?
                $(, varying: $varying:literal)?
                $(, since: $since:literal)?
                $(, platforms: [$($platform:literal),+ $(,)?])?
                $(,)?
            ;
        )*
    ) => {
        /// The metadata of each tool declared in this module, one constant per
        /// tool, so that a handler needing its own row — for the version it
        /// requires, or for its name in an error — names it rather than
        /// searching for it.
        pub mod metas {
            #![allow(non_upper_case_globals)]

            $(
                #[doc = ::std::concat!($($summary),*)]
                pub const $name: $crate::meta::ToolMeta = {
                    #[allow(unused_mut, unused_assignments)]
                    let mut severing = false;
                    $( severing = $severing; )?
                    #[allow(unused_mut, unused_assignments)]
                    let mut confirm = severing;
                    $( confirm = $confirm || severing; )?
                    #[allow(unused_mut, unused_assignments)]
                    let mut severs_local_node = false;
                    $( severs_local_node = $severs_local; )?
                    #[allow(unused_mut, unused_assignments)]
                    let mut idempotent = false;
                    $( idempotent = $idempotent; )?
                    #[allow(unused_mut, unused_assignments)]
                    let mut varying_tier = false;
                    $( varying_tier = $varying; )?
                    #[allow(unused_mut, unused_assignments)]
                    let mut since: ::std::option::Option<&'static str> =
                        ::std::option::Option::None;
                    $( since = ::std::option::Option::Some($since); )?
                    #[allow(unused_mut, unused_assignments)]
                    let mut platforms: ::std::option::Option<&'static [&'static str]> =
                        ::std::option::Option::None;
                    $( platforms = ::std::option::Option::Some(&[$($platform),+]); )?

                    $crate::meta::ToolMeta {
                        name: ::std::stringify!($name),
                        toolset: $crate::meta::Toolset::$toolset,
                        tier: $crate::meta::Tier::$tier,
                        summary: ::std::concat!($($summary),*).trim_ascii(),
                        self_severing: severing,
                        severs_local_node,
                        requires_confirmation: confirm,
                        idempotent,
                        varying_tier,
                        min_version: since,
                        platforms,
                    }
                };
            )*
        }

        /// Every tool declared in this module, in declaration order.
        pub fn entries() -> ::std::vec::Vec<$crate::registry::ToolEntry> {
            ::std::vec![
                $(
                    {
                        // Each tool gets its own shim so that the table can
                        // hold plain function pointers rather than boxed
                        // closures over twenty different parameter types.
                        fn invoke(
                            ctx: ::std::sync::Arc<$crate::context::ToolContext>,
                            args: ::rmcp::model::JsonObject,
                        ) -> ::tailscale_cli::BoxFuture<
                            'static,
                            $crate::error::ToolResult<::serde_json::Value>,
                        > {
                            ::std::boxed::Box::pin(async move {
                                let params: $params = $crate::registry::parse_params(
                                    ::std::stringify!($name),
                                    args,
                                )?;
                                $handler(&ctx, params).await
                            })
                        }

                        $crate::registry::ToolEntry {
                            meta: self::metas::$name,
                            schema: || {
                                ::rmcp::handler::server::tool::schema_for_input::<$params>()
                            },
                            invoke,
                        }
                    },
                )*
            ]
        }
    };
}

/// Turn a call's arguments into a handler's parameter type.
///
/// A shape mismatch is `invalid_args` rather than a protocol error: the caller
/// is a model, and a structured answer naming the problem is something it can
/// act on, whereas a JSON-RPC error is opaque to it.
///
/// An argument the tool does not take is refused rather than dropped, at any
/// depth the schema names: dropped, a misspelling becomes a call that reports
/// success for something it did not do (Q172). An object the schema gives no
/// property list, such as a document in Tailscale's own shape, stays open.
pub fn parse_params<T: serde::de::DeserializeOwned + schemars::JsonSchema + 'static>(
    tool: &str,
    args: JsonObject,
) -> ToolResult<T> {
    let args = Value::Object(args);
    if let Ok(schema) = rmcp::handler::server::tool::schema_for_input::<T>() {
        let root = Value::Object(schema.as_ref().clone());
        if let Some((path, accepted)) = unknown_argument_in(&root, &args) {
            let takes = if accepted.is_empty() {
                "no arguments".to_owned()
            } else {
                accepted.join(", ")
            };
            return Err(ToolError::invalid_args(format!(
                "`{tool}` takes no argument `{path}`; it takes {takes}"
            )));
        }
    }
    serde_json::from_value(args).map_err(|e| {
        ToolError::invalid_args(format!("`{tool}` was called with unusable arguments: {e}"))
    })
}

/// The first argument in `args` that a tool's input `schema` has no property
/// for, by path, with the names accepted at that level.
pub(crate) fn unknown_argument_in(schema: &Value, args: &Value) -> Option<(String, Vec<String>)> {
    unknown_argument(schema, schema, args, "")
}

/// The first argument in `value` that `schema` has no property for, with the
/// names it does have there.
fn unknown_argument(
    root: &Value,
    schema: &Value,
    value: &Value,
    path: &str,
) -> Option<(String, Vec<String>)> {
    let branches = branches(root, schema);
    match value {
        Value::Object(fields) => {
            let open = branches.iter().any(|b| {
                b.get("additionalProperties")
                    .is_some_and(|extra| extra != &Value::Bool(false))
            });
            let properties: Vec<&JsonObject> = branches
                .iter()
                .filter_map(|b| b.get("properties")?.as_object())
                .collect();
            // A struct with no fields declares an object and lists nothing; a
            // free-form value declares neither.
            let declared = !properties.is_empty() || branches.iter().any(|b| is_object_type(b));
            if open || !declared {
                return None;
            }
            fields.iter().find_map(|(key, field)| {
                let here = if path.is_empty() {
                    key.clone()
                } else {
                    format!("{path}.{key}")
                };
                match properties.iter().find_map(|p| p.get(key)) {
                    Some(property) => unknown_argument(root, property, field, &here),
                    None => {
                        let mut accepted: Vec<String> = properties
                            .iter()
                            .flat_map(|p| p.keys().map(|k| format!("`{k}`")))
                            .collect();
                        accepted.sort();
                        accepted.dedup();
                        Some((here, accepted))
                    }
                }
            })
        }
        Value::Array(items) => {
            let item = branches.iter().find_map(|b| b.get("items"))?;
            items.iter().enumerate().find_map(|(i, element)| {
                unknown_argument(root, item, element, &format!("{path}[{i}]"))
            })
        }
        _ => None,
    }
}

/// Whether a schema says its value is an object.
fn is_object_type(schema: &JsonObject) -> bool {
    match schema.get("type") {
        Some(Value::String(kind)) => kind == "object",
        Some(Value::Array(kinds)) => kinds.iter().any(|kind| kind == "object"),
        _ => false,
    }
}

/// A schema and every alternative it is made of, with references followed.
fn branches<'a>(root: &'a Value, schema: &'a Value) -> Vec<&'a JsonObject> {
    let Some(object) = schema.as_object() else {
        return Vec::new();
    };
    if let Some(target) = object.get("$ref").and_then(Value::as_str) {
        let resolved = target
            .strip_prefix("#/")
            .and_then(|pointer| root.pointer(&format!("/{pointer}")));
        return resolved.map_or_else(Vec::new, |r| branches(root, r));
    }
    let mut all = vec![object];
    for key in ["allOf", "anyOf", "oneOf"] {
        if let Some(alternatives) = object.get(key).and_then(Value::as_array) {
            all.extend(alternatives.iter().flat_map(|a| branches(root, a)));
        }
    }
    all
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::meta::{Tier, Toolset};
    use schemars::JsonSchema;
    use serde::Deserialize;
    use std::collections::BTreeSet;

    #[derive(Debug, Deserialize, JsonSchema)]
    struct Empty {}

    /// A parameter struct exercising both naming conventions in one tool, as
    /// ADR-0004 requires: what we own is snake_case, what Tailscale owns keeps
    /// the shape Tailscale gave it.
    // The fields exist to be reflected into a schema, not to be read.
    #[allow(dead_code)]
    #[derive(Debug, Deserialize, JsonSchema)]
    struct MixedParams {
        /// Ours: how the result should be shaped.
        include_peers: bool,
        /// Theirs: forwarded to the control plane unchanged.
        #[serde(rename = "keyExpiryDisabled")]
        key_expiry_disabled: bool,
    }

    async fn ok_handler<P>(_ctx: &ToolContext, _params: P) -> ToolResult<Value> {
        Ok(Value::Null)
    }

    mod declared {
        use super::*;

        crate::tools! {
            /// Report the state of the local node.
            tailscale_status => Empty, super::ok_handler,
                toolset: LocalStatus, tier: Read, idempotent: true;

            /// Disconnect this node from the tailnet.
            tailscale_down => Empty, super::ok_handler,
                toolset: LocalPrefs, tier: Destructive, severing: true;

            /// Delete a device from the tailnet.
            tailnet_device_delete => MixedParams, super::ok_handler,
                toolset: TailnetDevices, tier: Destructive, confirm: true, since: "1.60";
        }
    }

    fn registry() -> Registry {
        Registry::new(declared::entries()).expect("a well-formed table")
    }

    fn open_gate() -> Gate {
        Gate::unchecked(
            Toolset::ALL.iter().copied().collect(),
            Tier::Destructive,
            BTreeSet::new(),
        )
    }

    #[test]
    fn a_declaration_produces_its_metadata_row() {
        let registry = registry();
        assert_eq!(registry.len(), 3);

        let status = registry.get("tailscale_status").expect("declared");
        assert_eq!(status.meta.toolset, Toolset::LocalStatus);
        assert_eq!(status.meta.tier, Tier::Read);
        assert_eq!(status.meta.summary, "Report the state of the local node.");
        assert!(status.meta.idempotent);
        assert!(!status.meta.requires_confirmation);
        assert_eq!(status.meta.min_version, None);
    }

    #[test]
    fn a_tool_that_severs_the_local_node_must_offer_somewhere_to_say_so() {
        mod without {
            use crate::registry::tests::Empty;

            crate::tools! {
                /// Delete a device, claiming to sever without asking.
                tailnet_device_delete => Empty, super::ok_handler,
                    toolset: TailnetDevices, tier: Destructive, severs_local: true;
            }
        }
        // The flag alone would be a claim the schema does not keep: `Empty`
        // has no `confirm`, so nothing a caller sends could confirm anything.
        assert_eq!(
            Registry::new(without::entries()).expect_err("no `confirm` to give"),
            RegistryError::SeversLocalWithoutConfirmParameter("tailnet_device_delete")
        );
    }

    #[test]
    fn optional_settings_land_where_they_are_given() {
        let delete = registry()
            .get("tailnet_device_delete")
            .expect("declared")
            .meta;
        assert!(delete.requires_confirmation);
        assert!(!delete.self_severing);
        assert_eq!(delete.min_version, Some("1.60"));
    }

    #[test]
    fn a_duplicate_name_is_refused_at_construction() {
        let mut entries = declared::entries();
        entries.push(entries[0].clone());
        assert_eq!(
            Registry::new(entries).err(),
            Some(RegistryError::DuplicateName("tailscale_status"))
        );
    }

    #[test]
    fn a_name_the_protocol_or_our_prefixes_would_not_allow_is_refused() {
        let too_long: &'static str =
            Box::leak(format!("tailscale_{}", "a".repeat(MAX_NAME_LEN)).into_boxed_str());
        for (name, reason) in [
            ("", "empty"),
            (too_long, "128 characters"),
            ("tailscale_ping peer", "outside [A-Za-z0-9_.-]"),
            ("ping", "does not begin"),
            ("tailscaleping", "does not begin"),
        ] {
            match validate_name(name) {
                Err(RegistryError::BadName {
                    name: refused,
                    reason: said,
                }) => {
                    assert_eq!(refused, name);
                    assert!(said.contains(reason), "`{name}` {said}");
                }
                other => panic!("`{name}` should be refused: {other:?}"),
            }
        }
        assert_eq!(validate_name("tailnet_split-dns.v2"), Ok(()));
    }

    #[test]
    fn a_confirming_tool_refuses_until_the_caller_says_so() {
        let registry = registry();
        let gate = open_gate();

        let err = registry
            .resolve("tailscale_down", JsonObject::new(), &gate)
            .expect_err("must not run unconfirmed");
        assert_eq!(err.code, crate::error::ErrorCode::ConfirmationRequired);
        assert!(err.message.contains("disconnect"), "{}", err.message);

        let mut args = JsonObject::new();
        args.insert(CONFIRM_PARAM.to_owned(), Value::Bool(true));
        let (entry, rest) = registry
            .resolve("tailscale_down", args, &gate)
            .expect("a confirmed call runs");
        assert_eq!(entry.meta.name, "tailscale_down");
        assert!(
            !rest.contains_key(CONFIRM_PARAM),
            "the handler should not see the flag"
        );
    }

    #[test]
    fn a_confirmation_flag_of_the_wrong_shape_is_a_bad_argument() {
        let mut args = JsonObject::new();
        args.insert(CONFIRM_PARAM.to_owned(), Value::String("yes".to_owned()));
        let err = registry()
            .resolve("tailscale_down", args, &open_gate())
            .expect_err("`yes` is not a boolean");
        assert_eq!(err.code, crate::error::ErrorCode::InvalidArgs);
    }
}