use std::{
collections::{BTreeSet, HashMap},
sync::{
Arc, RwLock,
atomic::{AtomicU64, Ordering},
},
};
#[derive(Debug, Default)]
pub(crate) struct AgentTools {
pub(crate) hidden: BTreeSet<String>,
#[cfg(feature = "mcp")]
pub(crate) mcp_servers: Vec<String>,
pub(crate) host_tools: Vec<String>,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
struct AgentOwner(u64);
#[derive(Debug)]
struct Entry {
owner: AgentOwner,
tools: Arc<AgentTools>,
}
#[derive(Debug, Default)]
pub(crate) struct AgentRegistry {
agents: RwLock<HashMap<String, Entry>>,
next_owner: AtomicU64,
}
impl AgentRegistry {
pub(crate) fn of(&self, agent_id: &str) -> Option<Arc<AgentTools>> {
self.agents
.read()
.expect("agent registry poisoned")
.get(agent_id)
.map(|entry| Arc::clone(&entry.tools))
}
fn new_owner(&self) -> AgentOwner {
AgentOwner(self.next_owner.fetch_add(1, Ordering::Relaxed))
}
pub(crate) fn adopt(self: &Arc<Self>, parent: &str, child: &str) -> Option<AgentRow> {
let tools = self.of(parent)?;
let owner = self.new_owner();
self.agents
.write()
.expect("agent registry poisoned")
.insert(child.to_string(), Entry { owner, tools });
Some(AgentRow {
registry: Arc::clone(self),
owner,
agent_id: child.to_string(),
})
}
pub(crate) fn record_for_workspace(
self: &Arc<Self>,
agent_id: &str,
tools: AgentTools,
) -> AgentRow {
let owner = self.new_owner();
self.record(owner, agent_id, Arc::new(tools));
AgentRow {
registry: Arc::clone(self),
owner,
agent_id: agent_id.to_string(),
}
}
fn record(&self, owner: AgentOwner, agent_id: &str, tools: Arc<AgentTools>) {
self.agents
.write()
.expect("agent registry poisoned")
.insert(agent_id.to_string(), Entry { owner, tools });
}
fn forget_if_owned(&self, agent_id: &str, owner: AgentOwner) {
let mut agents = self
.agents
.write()
.unwrap_or_else(std::sync::PoisonError::into_inner);
if agents
.get(agent_id)
.is_some_and(|entry| entry.owner == owner)
{
agents.remove(agent_id);
}
}
}
#[derive(Debug)]
#[must_use = "dropping the hold takes the agent's row off the ledger"]
pub(crate) struct AgentRow {
registry: Arc<AgentRegistry>,
owner: AgentOwner,
agent_id: String,
}
impl Drop for AgentRow {
fn drop(&mut self) {
self.registry.forget_if_owned(&self.agent_id, self.owner);
}
}
#[derive(Debug)]
pub(crate) struct ForeignToolGuard {
agents: Arc<AgentRegistry>,
claims: crate::runtime::ToolClaims,
}
impl ForeignToolGuard {
pub(crate) fn new(agents: Arc<AgentRegistry>, claims: crate::runtime::ToolClaims) -> Self {
Self { agents, claims }
}
}
#[async_trait::async_trait]
impl crate::hooks::Interceptor for ForeignToolGuard {
fn name(&self) -> &str {
"basis tool ownership"
}
async fn intercept(
&self,
call: &crate::hooks::HookRequest,
) -> Result<crate::hooks::HookOutcome, crate::hooks::InterceptorError> {
let tools = self.agents.of(&call.agent_id);
#[cfg(feature = "mcp")]
if let Some((server, _)) = mentra::mcp::parse_mcp_tool_name(&call.tool_name) {
let Some(tools) = tools.as_ref() else {
return Ok(crate::hooks::HookOutcome::Allow);
};
return Ok(if tools.mcp_servers.iter().any(|own| own == server) {
crate::hooks::HookOutcome::Allow
} else {
crate::hooks::HookOutcome::Deny(format!(
"'{}' belongs to the MCP server '{server}', which this workspace did not \
configure",
call.tool_name
))
});
}
let owns = tools
.as_ref()
.is_some_and(|tools| tools.host_tools.contains(&call.tool_name));
if owns || !self.claims.holds_native(&call.tool_name) {
return Ok(crate::hooks::HookOutcome::Allow);
}
Ok(crate::hooks::HookOutcome::Deny(format!(
"'{}' is a native tool another open of this workspace supplied, and this one did not",
call.tool_name
)))
}
}
#[cfg(test)]
mod tests {
use super::*;
fn tools(hidden: &[&str]) -> AgentTools {
AgentTools {
hidden: hidden.iter().map(|name| (*name).to_string()).collect(),
#[cfg(feature = "mcp")]
mcp_servers: Vec::new(),
host_tools: Vec::new(),
}
}
#[cfg(feature = "mcp")]
#[tokio::test]
async fn the_guard_judges_a_bridged_name_by_the_callers_own_servers() {
use crate::hooks::{HookEvent, HookOutcome, HookRequest, Interceptor};
let registry = Arc::new(AgentRegistry::default());
let _owner = registry.record_for_workspace(
"owner",
AgentTools {
hidden: BTreeSet::new(),
mcp_servers: vec!["prod-db".to_string()],
host_tools: Vec::new(),
},
);
let _stranger = registry.record_for_workspace("stranger", tools(&[]));
let guard = ForeignToolGuard::new(registry, crate::runtime::ToolClaims::default());
let call = |agent: &str, tool: &str| HookRequest {
hook_schema: 1,
event: HookEvent::PreToolUse,
workspace: std::path::PathBuf::from("/repo"),
agent_id: agent.to_string(),
tool_call_id: "call-0".to_string(),
tool_name: tool.to_string(),
input: serde_json::json!({}),
output: None,
is_error: None,
};
assert!(matches!(
guard
.intercept(&call("owner", "mcp__prod-db__query"))
.await
.expect("decides"),
HookOutcome::Allow
));
assert!(
matches!(
guard
.intercept(&call("stranger", "mcp__prod-db__query"))
.await
.expect("decides"),
HookOutcome::Deny(_)
),
"the open that configured no servers must not reach the other's"
);
assert!(
matches!(
guard
.intercept(&call("stranger", "read"))
.await
.expect("decides"),
HookOutcome::Allow
),
"a name that is not a bridged tool's is none of this guard's business"
);
assert!(
matches!(
guard
.intercept(&call("host-session", "mcp__prod-db__query"))
.await
.expect("decides"),
HookOutcome::Allow
),
"a session basis never minted carries no workspace's guards at all"
);
}
#[test]
fn a_row_leaves_with_the_run_that_holds_it() {
let registry = Arc::new(AgentRegistry::default());
let one = registry.record_for_workspace("agent-1", tools(&["mcp__prod-db__query"]));
let two = registry.record_for_workspace("agent-2", tools(&["mcp__prod-db__query"]));
assert!(
registry.of("agent-1").is_some() && registry.of("agent-2").is_some(),
"a session whose workspace has gone is still a session, and still has to be \
judged by what its own open configured"
);
drop(one);
assert!(
registry.of("agent-1").is_none() && registry.of("agent-2").is_some(),
"each row leaves with the last hold on it and takes no other with it"
);
drop(two);
assert!(
registry.of("agent-2").is_none(),
"and nothing outlives the thing that needed it"
);
}
fn hides(registry: &AgentRegistry, agent_id: &str) -> Vec<String> {
registry
.of(agent_id)
.expect("a recorded agent")
.hidden
.iter()
.cloned()
.collect()
}
#[test]
fn a_hold_only_releases_the_row_it_still_owns() {
let registry = Arc::new(AgentRegistry::default());
let a_moved = registry.record_for_workspace("moved", tools(&["mcp__a-only__query"]));
let a_own = registry.record_for_workspace("a's own", tools(&["mcp__a-only__query"]));
let b_moved = registry.record_for_workspace("moved", tools(&["mcp__b-only__query"]));
let a_moved_again = registry.record_for_workspace("moved", tools(&["mcp__a-only__query"]));
drop(b_moved);
drop(a_moved);
assert_eq!(
hides(®istry, "moved"),
["mcp__a-only__query"],
"neither the sibling that lost this agent nor the stale hold from before it \
came back may erase the row the live run re-recorded: a missing row is a \
guard that allows and a child that inherits no hides"
);
drop(a_moved_again);
drop(a_own);
assert!(
registry.of("moved").is_none() && registry.of("a's own").is_none(),
"and the holds that do own their rows still release them"
);
}
#[test]
fn an_agent_a_sibling_took_over_outlives_the_run_that_minted_it() {
let registry = Arc::new(AgentRegistry::default());
let a_moved = registry.record_for_workspace("moved", tools(&["mcp__a-only__query"]));
let a_stayed = registry.record_for_workspace("stayed", tools(&["mcp__a-only__query"]));
let b_moved = registry.record_for_workspace("moved", tools(&["mcp__b-only__query"]));
drop(a_moved);
drop(a_stayed);
assert_eq!(
hides(®istry, "moved"),
["mcp__b-only__query"],
"the row belongs to whoever wrote it last, and that is who is running it"
);
assert!(
registry.of("stayed").is_none(),
"declining to release one agent must not hold back the others"
);
drop(b_moved);
assert!(
registry.of("moved").is_none(),
"and the row still leaves with the hold that does own it"
);
}
#[test]
fn a_delegated_child_answers_for_its_parents_workspace_until_it_returns() {
let registry = Arc::new(AgentRegistry::default());
let _parent = registry.record_for_workspace("parent", tools(&["mcp__prod-db__query"]));
let adopted = registry
.adopt("parent", "child")
.expect("the parent is known");
assert_eq!(
registry
.of("child")
.expect("the child inherits")
.hidden
.iter()
.map(String::as_str)
.collect::<Vec<_>>(),
["mcp__prod-db__query"],
"a child inherits its parent's audience, so it inherits its parent's denials"
);
drop(adopted);
assert!(registry.of("child").is_none());
assert!(
registry.of("parent").is_some(),
"the parent is still running"
);
assert!(
registry.adopt("stranger", "grandchild").is_none(),
"an agent basis did not make has nothing to hand down"
);
}
}