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;
pub const CONFIRM_PARAM: &str = "confirm";
pub const MAX_NAME_LEN: usize = 128;
pub type InvokeFn = fn(Arc<ToolContext>, JsonObject) -> BoxFuture<'static, ToolResult<Value>>;
#[derive(Clone)]
pub struct ToolEntry {
pub meta: ToolMeta,
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 {
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),
),
)
}
}
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.",
}),
);
}
}
#[derive(Debug, Clone)]
pub struct Registry {
entries: Vec<ToolEntry>,
by_name: BTreeMap<&'static str, usize>,
}
impl Registry {
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 }),
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 })
}
pub fn metas(&self) -> Vec<ToolMeta> {
self.entries.iter().map(|e| e.meta).collect()
}
pub fn entries(&self) -> &[ToolEntry] {
&self.entries
}
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()
}
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
}
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))
}
}
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`"
}
}
#[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 },
}
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(())
}
#[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),+ $(,)?])?
$(,)?
;
)*
) => {
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,
}
};
)*
}
pub fn entries() -> ::std::vec::Vec<$crate::registry::ToolEntry> {
::std::vec![
$(
{
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,
}
},
)*
]
}
};
}
pub fn parse_params<T: serde::de::DeserializeOwned>(tool: &str, args: JsonObject) -> ToolResult<T> {
serde_json::from_value(Value::Object(args)).map_err(|e| {
ToolError::invalid_args(format!("`{tool}` was called with unusable arguments: {e}"))
})
}
#[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 {}
#[allow(dead_code)]
#[derive(Debug, Deserialize, JsonSchema)]
struct MixedParams {
include_peers: bool,
#[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! {
tailscale_status => Empty, super::ok_handler,
toolset: LocalStatus, tier: Read, idempotent: true;
tailscale_down => Empty, super::ok_handler,
toolset: LocalPrefs, tier: Destructive, severing: true;
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! {
tailnet_device_delete => Empty, super::ok_handler,
toolset: TailnetDevices, tier: Destructive, severs_local: true;
}
}
assert_eq!(
Registry::new(without::entries()).expect_err("no `confirm` to give"),
RegistryError::SeversLocalWithoutConfirmParameter("tailnet_device_delete")
);
}
#[test]
fn a_self_severing_tool_requires_confirmation_without_being_told_twice() {
let down = registry().get("tailscale_down").expect("declared").meta;
assert!(down.self_severing);
assert!(
down.requires_confirmation,
"severing must imply confirmation"
);
}
#[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 every_tool_has_exactly_one_row_and_a_usable_name() {
let registry = registry();
let mut names: Vec<&str> = registry.entries().iter().map(|e| e.meta.name).collect();
let before = names.len();
names.sort_unstable();
names.dedup();
assert_eq!(before, names.len(), "a tool appears twice");
for name in names {
assert!(!name.is_empty());
assert!(name.len() <= MAX_NAME_LEN, "{name} is too long");
assert!(
name.bytes()
.all(|b| b.is_ascii_alphanumeric() || b == b'_' || b == b'-' || b == b'.'),
"{name} uses a character the protocol does not allow"
);
assert!(
name.starts_with("tailscale_") || name.starts_with("tailnet_"),
"{name} does not say which surface it acts on"
);
}
}
#[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 annotations_are_derived_from_the_tier_not_declared_beside_it() {
let registry = registry();
let status = registry
.get("tailscale_status")
.expect("declared")
.describe()
.expect("a valid schema");
let a = status.annotations.expect("annotations are set");
assert_eq!(a.read_only_hint, Some(true));
assert_eq!(a.destructive_hint, Some(false));
assert_eq!(a.idempotent_hint, Some(true));
assert_eq!(a.open_world_hint, Some(true));
let down = registry
.get("tailscale_down")
.expect("declared")
.describe()
.expect("a valid schema");
let a = down.annotations.expect("annotations are set");
assert_eq!(a.read_only_hint, Some(false));
assert_eq!(a.destructive_hint, Some(true));
}
#[test]
fn both_naming_conventions_survive_schema_generation() {
let tool = registry()
.get("tailnet_device_delete")
.expect("declared")
.describe()
.expect("a valid schema");
let properties = tool
.input_schema
.get("properties")
.and_then(Value::as_object)
.expect("an object schema with properties");
assert!(
properties.contains_key("include_peers"),
"server-owned parameters stay snake_case: {properties:?}"
);
assert!(
properties.contains_key("keyExpiryDisabled"),
"control-plane fields keep their own shape: {properties:?}"
);
}
#[test]
fn the_confirmation_flag_appears_only_where_it_is_required() {
let registry = registry();
let confirming = registry
.get("tailscale_down")
.expect("declared")
.describe()
.expect("a valid schema");
let properties = confirming
.input_schema
.get("properties")
.and_then(Value::as_object)
.expect("properties");
assert!(properties.contains_key(CONFIRM_PARAM));
let plain = registry
.get("tailscale_status")
.expect("declared")
.describe()
.expect("a valid schema");
let properties = plain
.input_schema
.get("properties")
.and_then(Value::as_object);
assert!(properties.is_none_or(|p| !p.contains_key(CONFIRM_PARAM)));
}
#[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);
}
#[test]
fn the_router_is_the_table_filtered_by_the_gate() {
let registry = registry();
let gate = Gate::unchecked(
BTreeSet::from([Toolset::LocalStatus, Toolset::LocalPrefs]),
Tier::Read,
BTreeSet::new(),
);
let visible: Vec<&str> = registry
.visible(&gate)
.iter()
.map(|e| e.meta.name)
.collect();
assert_eq!(visible, ["tailscale_status"]);
let err = registry
.resolve("tailscale_down", JsonObject::new(), &gate)
.expect_err("hidden tools do not run");
assert_eq!(err.code, crate::error::ErrorCode::NotPermitted);
assert!(
err.hint
.as_deref()
.is_some_and(|h| h.contains("--allow-destructive")),
"{err:?}"
);
}
#[test]
fn an_unknown_name_is_not_found() {
let err = registry()
.resolve("tailscale_nonesuch", JsonObject::new(), &open_gate())
.expect_err("no such tool");
assert_eq!(err.code, crate::error::ErrorCode::NotFound);
}
#[test]
fn visible_tools_are_listed_in_a_stable_order() {
let registry = registry();
let names: Vec<&str> = registry
.visible(&open_gate())
.iter()
.map(|e| e.meta.name)
.collect();
let mut sorted = names.clone();
sorted.sort_unstable();
assert_eq!(names, sorted);
}
}