use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, RwLock, RwLockReadGuard, RwLockWriteGuard};
use mentra::{
skill_root_key,
tool::{AudienceToolRegistration, PreparedTool, ToolAudience, ToolNameCollision},
};
use crate::tools::declared::DeclaredToolSpec;
use super::Runtime;
#[cfg(feature = "mcp")]
#[derive(Debug)]
pub(super) struct McpClaim {
root: PathBuf,
tools: Vec<String>,
}
#[derive(Debug)]
pub(super) struct ToolNameClaim {
root: PathBuf,
holders: usize,
program: ClaimedProgram,
registration: Option<AudienceToolRegistration>,
}
#[derive(Debug)]
enum ClaimedProgram {
Declared {
spec: Box<DeclaredToolSpec>,
supplied_holders: usize,
},
Native,
}
#[derive(Debug, Clone, Default)]
pub(crate) struct ToolClaims(Arc<RwLock<HashMap<String, ToolNameClaim>>>);
impl ToolClaims {
fn read(&self) -> RwLockReadGuard<'_, HashMap<String, ToolNameClaim>> {
self.0.read().expect("tool claim map poisoned")
}
fn lock(&self) -> RwLockWriteGuard<'_, HashMap<String, ToolNameClaim>> {
self.0.write().expect("tool claim map poisoned")
}
pub(crate) fn holds_native(&self, name: &str) -> bool {
self.read()
.get(name)
.is_some_and(|claim| matches!(claim.program, ClaimedProgram::Native))
}
pub(crate) fn foreign_native_on(
&self,
root: &Path,
own: &[String],
) -> std::collections::BTreeSet<String> {
self.read()
.iter()
.filter(|(name, claim)| {
claim.root == root
&& matches!(claim.program, ClaimedProgram::Native)
&& !own.iter().any(|mine| mine == *name)
})
.map(|(name, _)| name.clone())
.collect()
}
}
#[cfg(feature = "mcp")]
impl McpClaim {
fn new(root: &Path) -> Self {
Self {
root: root.to_path_buf(),
tools: Vec::new(),
}
}
}
#[derive(Debug)]
#[must_use = "a claimed name with nothing registered under it is a tool the model cannot call"]
pub(crate) struct ToolNamePermit {
name: String,
root: PathBuf,
}
impl ToolNamePermit {
fn issue(
claims: &mut HashMap<String, ToolNameClaim>,
name: &str,
root: &Path,
program: ClaimedProgram,
) -> Self {
claims.insert(
name.to_string(),
ToolNameClaim {
root: root.to_path_buf(),
holders: 1,
program,
registration: None,
},
);
Self {
name: name.to_string(),
root: root.to_path_buf(),
}
}
}
impl ToolNameClaim {
fn taken_elsewhere(&self) -> String {
let binding = match self.program {
ClaimedProgram::Declared { .. } => "declares a tool",
ClaimedProgram::Native => "supplies a native tool",
};
format!(
"the workspace at {} is open on this runtime and {binding} by that name",
self.root.display()
)
}
}
#[derive(Debug, Clone, Copy)]
pub(crate) enum DeclaredToolOrigin {
File,
Supplied,
}
impl Runtime {
#[cfg(feature = "mcp")]
pub(crate) fn claim_mcp_server(&self, name: &str, root: &Path) -> String {
let mut claims = self.mcp_claims.lock().expect("mcp claim map poisoned");
if !claims.contains_key(name) {
claims.insert(name.to_string(), McpClaim::new(root));
return name.to_string();
}
let mut effective = format!("{name}-{}", root_suffix(root));
let mut attempt = 2_u32;
while claims.contains_key(&effective) {
effective = format!("{name}-{}-{attempt}", root_suffix(root));
attempt += 1;
}
claims.insert(effective.clone(), McpClaim::new(root));
effective
}
#[cfg(feature = "mcp")]
pub(crate) fn record_bridged_tools(&self, name: &str, root: &Path, tools: Vec<String>) {
let mut claims = self.mcp_claims.lock().expect("mcp claim map poisoned");
if let Some(claim) = claims.get_mut(name)
&& claim.root == root
{
claim.tools = tools;
}
}
#[cfg(feature = "mcp")]
pub(crate) fn foreign_mcp_tools(&self, own: &[String]) -> std::collections::BTreeSet<String> {
let mine = |server: &str| own.iter().any(|owned| owned == server);
let mut foreign = std::collections::BTreeSet::new();
for (server, claim) in self
.mcp_claims
.lock()
.expect("mcp claim map poisoned")
.iter()
{
if mine(server) {
continue;
}
foreign.extend(claim.tools.iter().cloned());
}
for descriptor in self.mentra.tools() {
let name = &descriptor.provider.name;
if let Some((server, _)) = mentra::mcp::parse_mcp_tool_name(name)
&& !mine(server)
{
foreign.insert(name.clone());
}
}
foreign
}
#[cfg(feature = "mcp")]
pub(crate) fn release_mcp_claim(&self, name: &str, root: &Path) {
let mut claims = self.mcp_claims.lock().expect("mcp claim map poisoned");
if claims.get(name).is_some_and(|claim| claim.root == root) {
claims.remove(name);
}
}
pub(crate) fn claim_declared_tool(
&self,
root: &Path,
spec: &DeclaredToolSpec,
origin: DeclaredToolOrigin,
) -> Result<Option<ToolNamePermit>, String> {
let name = &spec.name;
let mut claims = self.tool_claims.lock();
match claims.get_mut(name) {
Some(claim) if claim.root != root => Err(claim.taken_elsewhere()),
Some(claim) => match &mut claim.program {
ClaimedProgram::Native => Err(
"another live open of this workspace supplied a native tool under that name"
.to_string(),
),
ClaimedProgram::Declared {
spec: held,
supplied_holders,
} if **held != *spec
&& (*supplied_holders > 0
|| matches!(origin, DeclaredToolOrigin::Supplied)) =>
{
Err(
"another live open of this workspace supplied different configuration \
under that name"
.to_string(),
)
}
ClaimedProgram::Declared {
supplied_holders, ..
} => {
if matches!(origin, DeclaredToolOrigin::Supplied) {
*supplied_holders += 1;
}
claim.holders += 1;
Ok(None)
}
},
None if self.registers_tool(name) => {
Err("this runtime already offers a tool by that name".to_string())
}
None => Ok(Some(ToolNamePermit::issue(
&mut claims,
name,
root,
ClaimedProgram::Declared {
spec: Box::new(spec.clone()),
supplied_holders: usize::from(matches!(origin, DeclaredToolOrigin::Supplied)),
},
))),
}
}
pub(crate) fn claim_native_tool(
&self,
root: &Path,
name: &str,
) -> Result<ToolNamePermit, String> {
let mut claims = self.tool_claims.lock();
match claims.get(name) {
Some(claim) if claim.root != root => Err(claim.taken_elsewhere()),
Some(_) => Err(
"this workspace is already open on this runtime with a tool by that name"
.to_string(),
),
None if self.registers_tool(name) => {
Err("this runtime already offers a tool by that name".to_string())
}
None => Ok(ToolNamePermit::issue(
&mut claims,
name,
root,
ClaimedProgram::Native,
)),
}
}
pub(crate) fn install_claimed_tool(
&self,
audience: &ToolAudience,
claim: ToolNamePermit,
prepared: PreparedTool,
) -> Result<(), String> {
if prepared.descriptor().provider.name != claim.name {
return Err(format!(
"its descriptor named '{}' when the name was claimed and '{}' when it was \
prepared for registration; a tool has to be the same tool both times",
claim.name,
prepared.descriptor().provider.name,
));
}
let mut claims = self.tool_claims.lock();
let Some(entry) = claims
.get_mut(&claim.name)
.filter(|entry| entry.root == claim.root)
else {
return Err(
"the claim on that name was released while this workspace was opening".to_string(),
);
};
let registration = self
.mentra
.try_register_prepared_tool_for_audience(audience.clone(), prepared)
.map_err(|collision: ToolNameCollision| {
format!(
"something registered a tool called '{}' on this runtime while this \
workspace was opening",
collision.name
)
})?;
entry.registration = Some(registration);
Ok(())
}
pub(crate) fn release_declared_tool(
&self,
name: &str,
root: &Path,
origin: DeclaredToolOrigin,
) {
self.release_tool_claim(name, root, matches!(origin, DeclaredToolOrigin::Supplied));
}
pub(crate) fn release_native_tool(&self, name: &str, root: &Path) {
self.release_tool_claim(name, root, false);
}
fn release_tool_claim(&self, name: &str, root: &Path, supplied: bool) {
let mut claims = self.tool_claims.lock();
let Some(claim) = claims.get_mut(name) else {
return;
};
if claim.root != root {
return;
}
claim.holders = claim.holders.saturating_sub(1);
if let ClaimedProgram::Declared {
supplied_holders, ..
} = &mut claim.program
&& supplied
{
*supplied_holders = supplied_holders.saturating_sub(1);
}
if claim.holders == 0 {
claims.remove(name);
}
}
pub(crate) fn register_skill_roots(
&self,
roots: &[PathBuf],
) -> Result<(), mentra::SkillLoadError> {
let mut holders = self
.skill_root_holders
.lock()
.expect("skill root holder map poisoned");
self.mentra.register_skills_dirs(roots)?;
for root in roots {
*holders.entry(skill_root_key(root)).or_insert(0) += 1;
}
Ok(())
}
pub(crate) fn release_skill_roots(&self, roots: &[PathBuf]) {
let mut holders = self
.skill_root_holders
.lock()
.expect("skill root holder map poisoned");
for root in roots {
let key = skill_root_key(root);
let Some(count) = holders.get_mut(&key) else {
continue;
};
*count = count.saturating_sub(1);
if *count == 0 {
holders.remove(&key);
self.mentra.unregister_skills_dir(&key);
}
}
}
pub(crate) fn foreign_native_tools(
&self,
root: &Path,
own: &[String],
) -> std::collections::BTreeSet<String> {
self.tool_claims.foreign_native_on(root, own)
}
pub(crate) fn tool_claims(&self) -> ToolClaims {
self.tool_claims.clone()
}
#[cfg(test)]
pub(crate) fn claimed_tool_descriptor(
&self,
name: &str,
) -> Option<mentra::tool::RuntimeToolDescriptor> {
self.tool_claims
.read()
.get(name)?
.registration
.as_ref()
.map(|registration| registration.descriptor().clone())
}
fn registers_tool(&self, name: &str) -> bool {
self.mentra
.tools()
.iter()
.any(|descriptor| descriptor.provider.name == name)
}
}
#[cfg(feature = "mcp")]
fn root_suffix(root: &Path) -> String {
const OFFSET: u64 = 0xcbf2_9ce4_8422_2325;
const PRIME: u64 = 0x0000_0100_0000_01b3;
let mut hash = OFFSET;
for byte in root.as_os_str().as_encoded_bytes() {
hash ^= u64::from(*byte);
hash = hash.wrapping_mul(PRIME);
}
format!("{:08x}", (hash >> 32) as u32 ^ hash as u32)
}
#[cfg(test)]
mod tests {
use super::*;
#[cfg(feature = "mcp")]
#[test]
fn a_taken_server_name_is_suffixed_and_a_released_one_is_free_again() {
use std::path::Path;
let runtime = Runtime::builder()
.with_base_url("http://127.0.0.1:1/v1")
.with_api_key("test-key")
.with_ephemeral_history()
.build()
.expect("builds");
let first = runtime.claim_mcp_server("fs", Path::new("/repo/one"));
let second = runtime.claim_mcp_server("fs", Path::new("/repo/two"));
let again = runtime.claim_mcp_server("fs", Path::new("/repo/two"));
assert_eq!(first, "fs", "the first claimant keeps the plain name");
assert_ne!(second, "fs", "the second must not collide in the registry");
assert!(second.starts_with("fs-"), "{second}");
assert_ne!(again, second, "every live claim is its own namespace");
runtime.release_mcp_claim("fs", Path::new("/repo/two"));
runtime.release_mcp_claim(&second, Path::new("/repo/one"));
assert_eq!(
runtime.claim_mcp_server("fs", Path::new("/repo/three")),
format!("fs-{}", root_suffix(Path::new("/repo/three"))),
"a name someone else holds stays held"
);
runtime.release_mcp_claim("fs", Path::new("/repo/one"));
assert_eq!(
runtime.claim_mcp_server("fs", Path::new("/repo/four")),
"fs",
"a released name is claimable plain again"
);
}
#[cfg(feature = "mcp")]
#[test]
fn a_bridged_tool_is_foreign_to_every_open_that_did_not_configure_its_server() {
use std::path::Path;
let runtime = Runtime::builder()
.with_base_url("http://127.0.0.1:1/v1")
.with_api_key("test-key")
.with_ephemeral_history()
.build()
.expect("builds");
let root = Path::new("/repo");
let server = runtime.claim_mcp_server("prod-db", root);
runtime.record_bridged_tools(&server, root, vec!["mcp__prod-db__query".to_string()]);
assert_eq!(
runtime
.foreign_mcp_tools(&[])
.into_iter()
.collect::<Vec<_>>(),
["mcp__prod-db__query"],
"the open that configured no servers must not be offered the other's"
);
assert!(
runtime
.foreign_mcp_tools(std::slice::from_ref(&server))
.is_empty(),
"and the open that configured it keeps it"
);
runtime.release_mcp_claim(&server, root);
assert!(runtime.foreign_mcp_tools(&[]).is_empty());
}
#[cfg(feature = "mcp")]
#[test]
fn a_global_tool_shaped_like_a_bridged_one_is_foreign_to_every_workspace() {
use mentra::tool::{
ParallelToolContext, RuntimeToolDescriptor, ToolDefinition, ToolExecutor, ToolResult,
};
use serde_json::{Value, json};
struct HostAdmin;
impl ToolDefinition for HostAdmin {
fn descriptor(&self) -> RuntimeToolDescriptor {
RuntimeToolDescriptor::builder("mcp__internal__admin")
.description("the host's own tool")
.input_schema(json!({"type": "object"}))
.build()
}
}
#[async_trait::async_trait]
impl ToolExecutor for HostAdmin {
async fn execute(&self, _ctx: ParallelToolContext, _input: Value) -> ToolResult {
Ok("administered".to_string())
}
}
let runtime = Runtime::builder()
.with_base_url("http://127.0.0.1:1/v1")
.with_api_key("test-key")
.with_ephemeral_history()
.with_tool(HostAdmin)
.build()
.expect("builds");
assert_eq!(
runtime
.foreign_mcp_tools(&[])
.into_iter()
.collect::<Vec<_>>(),
["mcp__internal__admin"],
"no workspace configured a server called `internal`"
);
}
}