use std::{
sync::Arc,
time::{Duration, Instant},
};
use async_trait::async_trait;
use polyc_agent::{ToolDecision, ToolExecutor};
use polyc_crypto::sensitive::Sensitive;
use polyc_llm::ToolSpec;
use rmcp::{
ClientCacheConfig, ClientLifecycleMode, ClientServiceExt,
model::{CallToolRequestParams, ProtocolVersion},
service::{RoleClient, RunningService},
transport::{
StreamableHttpClientTransport, streamable_http_client::StreamableHttpClientTransportConfig,
},
};
use crate::connection_pool::{ConnectionKey, ConnectionPool, PooledSession, SessionHandle};
use crate::connector_error::{
CallRetryPolicy, ConnectorErrorKind, call_error_message, classify_call_error, dial_error,
dial_failure_message, failure_json, result_message, success_json, transport_failure_message,
transport_source_is_auth,
};
pub const DEFAULT_CONNECT_TIMEOUT: Duration = Duration::from_secs(10);
pub const CONNECT_TIMEOUT: Duration = Duration::from_secs(5);
pub const CALLER_HEADER: &str = "x-polychrome-caller";
pub const CONNECTOR_TOOL_SEPARATOR: &str = "__";
#[must_use]
pub fn is_valid_connector_label(label: &str) -> bool {
!label.is_empty()
&& label
.chars()
.all(|c| c.is_ascii_lowercase() || c.is_ascii_digit() || c == '-')
}
#[derive(Debug, thiserror::Error)]
pub enum McpClientError {
#[error("MCP client initialise failed: {0}")]
Init(String),
#[error("MCP list_tools failed: {0}")]
ListTools(String),
#[error("MCP connect timed out after {0:?}")]
Timeout(Duration),
#[error("MCP connector rejected the credentials: {0}")]
AuthRejected(String),
#[error("invalid resource URI: {0}")]
InvalidResource(String),
#[error("trusted connector destination refused: {0}")]
BlockedDestination(String),
#[error(
"invalid connector label {0:?}: a label is a Kubernetes resource name \
(lowercase ASCII alphanumerics and `-`)"
)]
InvalidLabel(String),
#[error("bearer token audience {audience} does not match target resource {resource}")]
AudienceMismatch {
audience: String,
resource: String,
},
}
impl McpClientError {
#[must_use]
pub const fn kind(&self) -> ConnectorErrorKind {
match self {
Self::AuthRejected(_) | Self::AudienceMismatch { .. } => ConnectorErrorKind::Auth,
Self::InvalidLabel(_) | Self::InvalidResource(_) | Self::BlockedDestination(_) => {
ConnectorErrorKind::Config
}
Self::Init(_) | Self::ListTools(_) | Self::Timeout(_) => ConnectorErrorKind::Transport,
}
}
}
type Result<T> = std::result::Result<T, McpClientError>;
#[derive(Debug, Clone, Default)]
pub struct ApprovalPolicy {
pub connector: bool,
pub tools: Vec<String>,
}
impl ApprovalPolicy {
#[must_use]
pub const fn connector(connector: bool) -> Self {
Self {
connector,
tools: Vec::new(),
}
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ConnectorProvenance {
OperatorRegistered,
#[default]
SelfDeclared,
}
#[derive(Debug)]
pub struct ConnectOptions {
pub label: Option<String>,
pub bearer: Option<AudienceBoundToken>,
pub caller: Option<String>,
pub approval: ApprovalPolicy,
pub timeout: Option<Duration>,
pub source: SpecSource,
pub allowed_tools: Vec<String>,
}
impl Default for ConnectOptions {
fn default() -> Self {
Self {
label: None,
bearer: None,
caller: None,
approval: ApprovalPolicy::default(),
timeout: Some(DEFAULT_CONNECT_TIMEOUT),
source: SpecSource::List,
allowed_tools: Vec::new(),
}
}
}
fn canonical_resource(uri: &str) -> Result<String> {
let mut url = reqwest::Url::parse(uri)
.map_err(|e| McpClientError::InvalidResource(format!("{uri}: {e}")))?;
url.set_fragment(None);
Ok(String::from(url))
}
#[derive(Debug, Clone)]
pub struct AudienceBoundToken {
token: Sensitive<String>,
audience: String,
}
impl AudienceBoundToken {
pub fn new(token: impl Into<String>, resource_uri: &str) -> Result<Self> {
Ok(Self {
token: Sensitive::new(token.into()),
audience: canonical_resource(resource_uri)?,
})
}
#[must_use]
pub fn resource(&self) -> &str {
&self.audience
}
pub fn bearer_for(&self, target_uri: &str) -> Result<&str> {
let resource = canonical_resource(target_uri)?;
if resource == self.audience {
Ok(self.token.expose())
} else {
Err(McpClientError::AudienceMismatch {
audience: self.audience.clone(),
resource,
})
}
}
}
#[derive(Debug, Default)]
pub enum SpecSource {
#[default]
List,
Shipped(Vec<ToolSpec>),
}
#[must_use]
pub fn core_admits_builtin(name: &str, core: &[String]) -> bool {
!name.contains(CONNECTOR_TOOL_SEPARATOR) && core.iter().any(|n| n == name)
}
#[must_use]
pub fn builtin_admits(name: &str, allow: Option<&[String]>, core: &[String]) -> bool {
if core_admits_builtin(name, core) {
return true;
}
if crate::capability::management::is_management_builtin(name) {
allow.is_some_and(|allow| allow.iter().any(|n| n == name))
} else {
allow.is_none_or(|allow| allow.iter().any(|n| n == name))
}
}
pub(crate) async fn dial_service(
uri: Arc<str>,
auth_header: Option<String>,
caller: Option<String>,
trusted_transport: Option<reqwest::Client>,
) -> Result<RunningService<RoleClient, ()>> {
let http = match trusted_transport {
Some(http) => http,
None => reqwest::Client::builder()
.pool_max_idle_per_host(0)
.connect_timeout(CONNECT_TIMEOUT)
.redirect(reqwest::redirect::Policy::none())
.build()
.map_err(|e| McpClientError::Init(e.to_string()))?,
};
let mut config = match auth_header {
Some(token) => StreamableHttpClientTransportConfig::with_uri(uri).auth_header(token),
None => StreamableHttpClientTransportConfig::with_uri(uri),
};
if let Some(caller) = caller {
match http::HeaderValue::from_str(&caller) {
Ok(value) => {
config
.custom_headers
.insert(http::HeaderName::from_static(CALLER_HEADER), value);
}
Err(_) => {
tracing::warn!(
"caller identity is not a valid HTTP header value; dialing without it"
);
}
}
}
let transport = StreamableHttpClientTransport::with_client(http, config);
let service = ()
.serve_with_lifecycle(
transport,
ClientLifecycleMode::Discover {
preferred_versions: vec![ProtocolVersion::V_2026_07_28],
},
)
.await
.map_err(dial_error)?;
service
.peer()
.set_response_cache_config(catalog_cache_config())
.await;
Ok(service)
}
fn catalog_cache_config() -> ClientCacheConfig {
ClientCacheConfig::default().with_serve_stale_on_error(false)
}
fn gate_spec(mut spec: ToolSpec, approval: &ApprovalPolicy) -> ToolSpec {
let operator_gated = approval.tools.contains(&spec.name);
spec.needs_approval = approval.connector || spec.destructive || operator_gated;
spec
}
fn spec_from_rmcp_tool(t: rmcp::model::Tool) -> ToolSpec {
let destructive = t
.annotations
.as_ref()
.and_then(|a| a.destructive_hint)
.unwrap_or(false);
let read_only = t
.annotations
.as_ref()
.and_then(|a| a.read_only_hint)
.unwrap_or(false);
let open_world = t
.annotations
.as_ref()
.and_then(|a| a.open_world_hint)
.unwrap_or(true);
let mut spec = ToolSpec::new(
t.name.clone().into_owned(),
t.description
.map(std::borrow::Cow::into_owned)
.unwrap_or_default(),
serde_json::Value::Object((*t.input_schema).clone()),
);
spec.title = t.title;
spec.read_only = read_only;
spec.destructive = destructive;
spec.open_world = open_world;
spec
}
fn validate_prefix(label: Option<String>) -> Result<String> {
match label {
Some(label) if !is_valid_connector_label(&label) => {
Err(McpClientError::InvalidLabel(label))
}
Some(label) => Ok(format!("{label}{CONNECTOR_TOOL_SEPARATOR}")),
None => Ok(String::new()),
}
}
async fn list_specs(service: &RunningService<RoleClient, ()>) -> Result<Vec<ToolSpec>> {
Ok(service
.peer()
.list_all_tools()
.await
.map_err(|e| {
if transport_source_is_auth(&e) {
McpClientError::AuthRejected(e.to_string())
} else {
McpClientError::ListTools(e.to_string())
}
})?
.into_iter()
.map(spec_from_rmcp_tool)
.collect())
}
fn restrict_to_allowed(raw: Vec<ToolSpec>, allowed: &[String]) -> Vec<ToolSpec> {
if allowed.is_empty() {
return raw;
}
raw.into_iter()
.filter(|spec| allowed.iter().any(|name| name == &spec.name))
.collect()
}
fn finish_specs(
raw: Vec<ToolSpec>,
allowed: &[String],
approval: &ApprovalPolicy,
prefix: &str,
) -> Vec<ToolSpec> {
restrict_to_allowed(raw, allowed)
.into_iter()
.map(|spec| {
let mut spec = gate_spec(spec, approval);
spec.name = format!("{prefix}{}", spec.name);
spec
})
.collect()
}
pub struct McpToolSource {
session: SessionHandle,
specs: Vec<ToolSpec>,
needs_approval: bool,
provenance: ConnectorProvenance,
prefix: String,
retry: CallRetryPolicy,
}
impl std::fmt::Debug for McpToolSource {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("McpToolSource")
.field("session", &"<SessionHandle>")
.field(
"specs",
&self.specs.iter().map(|s| &s.name).collect::<Vec<_>>(),
)
.field("needs_approval", &self.needs_approval)
.field("provenance", &self.provenance)
.field("prefix", &self.prefix)
.field("retry", &self.retry)
.finish()
}
}
impl McpToolSource {
pub async fn connect(uri: impl Into<Arc<str>>, options: ConnectOptions) -> Result<Self> {
let timeout = options.timeout;
let dial = Self::dial(uri.into(), options, None);
match timeout {
Some(budget) => tokio::time::timeout(budget, dial)
.await
.map_err(|_elapsed| McpClientError::Timeout(budget))?,
None => dial.await,
}
}
pub async fn connect_with_transport(
uri: impl Into<Arc<str>>,
options: ConnectOptions,
transport: reqwest::Client,
) -> Result<Self> {
let timeout = options.timeout;
let dial = Self::dial(uri.into(), options, Some(transport));
match timeout {
Some(budget) => tokio::time::timeout(budget, dial)
.await
.map_err(|_elapsed| McpClientError::Timeout(budget))?,
None => dial.await,
}
}
async fn dial(
uri: Arc<str>,
options: ConnectOptions,
trusted_transport: Option<reqwest::Client>,
) -> Result<Self> {
let ConnectOptions {
label,
bearer,
caller,
approval,
source,
allowed_tools,
timeout: _,
} = options;
let prefix = validate_prefix(label)?;
let auth_header = match bearer {
Some(token) => Some(Sensitive::new(token.bearer_for(&uri)?.to_owned())),
None => None,
};
let service = dial_service(
uri,
auth_header.map(|h| h.expose().clone()),
caller,
trusted_transport,
)
.await?;
let raw = match source {
SpecSource::List => list_specs(&service).await?,
SpecSource::Shipped(specs) => specs,
};
let specs = finish_specs(raw, &allowed_tools, &approval, &prefix);
Ok(Self {
session: SessionHandle::Direct(Arc::new(service)),
specs,
needs_approval: approval.connector,
provenance: ConnectorProvenance::default(),
prefix,
retry: CallRetryPolicy::default(),
})
}
pub async fn pooled(
pool: ConnectionPool,
principal: impl Into<String>,
uri: impl Into<Arc<str>>,
options: ConnectOptions,
) -> Result<Self> {
let ConnectOptions {
label,
bearer,
caller,
approval,
timeout,
source,
allowed_tools,
} = options;
let uri: Arc<str> = uri.into();
let prefix = validate_prefix(label.clone())?;
let auth_header = match bearer {
Some(token) => Some(Sensitive::new(token.bearer_for(&uri)?.to_owned())),
None => None,
};
let key = ConnectionKey::new(canonical_resource(&uri)?, principal);
let pooled = PooledSession {
pool: pool.clone(),
key: key.clone(),
uri: Arc::clone(&uri),
auth_header: auth_header.clone(),
caller: caller.clone(),
connect_timeout: timeout,
label: label.unwrap_or_default(),
};
let raw = match source {
SpecSource::Shipped(specs) => specs,
SpecSource::List => {
let service = pool
.acquire(
&key,
&uri,
auth_header.as_ref().map(|h| h.expose().as_str()),
caller.as_deref(),
&pooled.label,
timeout,
)
.await?;
list_specs(&service).await?
}
};
let specs = finish_specs(raw, &allowed_tools, &approval, &prefix);
Ok(Self {
session: SessionHandle::Pooled(pooled),
specs,
needs_approval: approval.connector,
provenance: ConnectorProvenance::default(),
prefix,
retry: CallRetryPolicy::default(),
})
}
pub fn names(&self) -> impl Iterator<Item = &str> {
self.specs.iter().map(|s| s.name.as_str())
}
#[must_use]
pub const fn requires_approval(&self) -> bool {
self.needs_approval
}
#[must_use]
pub const fn operator_registered(mut self) -> Self {
self.provenance = ConnectorProvenance::OperatorRegistered;
self
}
#[must_use]
pub const fn provenance(&self) -> ConnectorProvenance {
self.provenance
}
#[must_use]
pub const fn with_call_retry_policy(mut self, retry: CallRetryPolicy) -> Self {
self.retry = retry;
self
}
pub fn merge_redeclared_specs(&mut self, redeclared: Vec<ToolSpec>) {
for new in redeclared {
if let Some(old) = self.specs.iter_mut().find(|s| s.name == new.name) {
old.description = new.description;
old.title = new.title;
old.schema_json = new.schema_json;
old.read_only = old.read_only && new.read_only;
old.cacheable_approval = old.cacheable_approval && new.cacheable_approval;
old.destructive = old.destructive || new.destructive;
old.open_world = old.open_world || new.open_world;
old.needs_approval = old.needs_approval || new.needs_approval;
} else {
let mut spec = new;
spec.needs_approval = spec.needs_approval || self.needs_approval;
self.specs.push(spec);
}
}
}
pub fn shutdown(&self) {
self.session.shutdown();
}
}
impl Drop for McpToolSource {
fn drop(&mut self) {
if let SessionHandle::Direct(service) = &self.session
&& Arc::strong_count(service) == 1
{
service.cancellation_token().cancel();
}
}
}
#[async_trait]
impl ToolExecutor for McpToolSource {
fn specs(&self) -> Vec<ToolSpec> {
self.specs.clone()
}
fn owns(&self, name: &str) -> bool {
self.specs.iter().any(|s| s.name == name)
}
fn required_capabilities(&self, name: &str) -> polyc_capability::CapabilitySet {
let Some(spec) = self.specs.iter().find(|s| s.name == name) else {
return polyc_capability::CapabilitySet::all();
};
let origin = match self.provenance {
ConnectorProvenance::OperatorRegistered => {
polyc_capability::ToolOrigin::RegisteredConnector
}
ConnectorProvenance::SelfDeclared => polyc_capability::ToolOrigin::Unknown,
};
polyc_capability::required_capabilities(polyc_capability::ToolProfile::for_spec(
spec, origin,
))
}
fn needs_approval(&self, name: &str) -> bool {
self.specs
.iter()
.find(|s| s.name == name)
.is_some_and(|s| s.needs_approval)
}
async fn execute(&self, name: &str, args_json: &str) -> String {
let arguments = serde_json::from_str::<serde_json::Value>(args_json)
.ok()
.and_then(|v| v.as_object().cloned());
let raw_name = name.strip_prefix(self.prefix.as_str()).unwrap_or(name);
let service = match self.session.acquire().await {
Ok(service) => service,
Err(err) => {
let kind = err.kind();
return failure_json(kind, &dial_failure_message(kind));
}
};
let started = Instant::now();
let mut attempt: u32 = 0;
let should_retry = |attempt: &mut u32| -> Option<Duration> {
let delay = polyc_agent::retry::backoff_delay(
*attempt,
self.retry.base_delay,
self.retry.max_delay,
polyc_agent::retry::Clock::jitter_frac(&polyc_agent::retry::RealClock),
);
(started.elapsed() + delay <= self.retry.budget).then(|| {
*attempt = attempt.saturating_add(1);
delay
})
};
loop {
let mut request = CallToolRequestParams::new(raw_name.to_owned());
request.arguments = arguments.clone();
match tokio::time::timeout(self.retry.call_timeout, service.call_tool(request)).await {
Ok(Ok(result)) if result.is_error == Some(true) => {
return failure_json(ConnectorErrorKind::Application, &result_message(&result));
}
Ok(Ok(result)) => return success_json(result),
Ok(Err(err)) => {
let kind = classify_call_error(&err);
if kind == ConnectorErrorKind::Transport
&& let Some(delay) = should_retry(&mut attempt)
{
tokio::time::sleep(delay).await;
continue;
}
if matches!(
kind,
ConnectorErrorKind::Transport | ConnectorErrorKind::Auth
) {
self.session.evict();
}
return failure_json(kind, &call_error_message(kind, &err));
}
Err(_elapsed) => {
if let Some(delay) = should_retry(&mut attempt) {
tokio::time::sleep(delay).await;
continue;
}
self.session.evict();
return failure_json(
ConnectorErrorKind::Transport,
transport_failure_message(),
);
}
}
}
}
}
pub struct CompositeRegistry {
sources: Vec<Arc<dyn ToolExecutor>>,
core: Vec<String>,
}
impl std::fmt::Debug for CompositeRegistry {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.debug_struct("CompositeRegistry")
.field("sources", &self.sources.len())
.field("core", &self.core)
.finish()
}
}
impl Default for CompositeRegistry {
fn default() -> Self {
Self::new()
}
}
impl CompositeRegistry {
#[must_use]
pub fn new() -> Self {
Self {
sources: Vec::new(),
core: Vec::new(),
}
}
#[must_use]
pub fn with_core(mut self, core: Vec<String>) -> Self {
self.core = core;
self
}
#[must_use]
pub fn with(mut self, executor: Arc<dyn ToolExecutor>) -> Self {
self.sources.push(executor);
self
}
pub fn push(&mut self, executor: Arc<dyn ToolExecutor>) {
self.sources.push(executor);
}
fn owner_of(&self, name: &str) -> Option<&Arc<dyn ToolExecutor>> {
self.sources.iter().find(|source| source.owns(name))
}
}
#[async_trait]
impl ToolExecutor for CompositeRegistry {
fn specs(&self) -> Vec<ToolSpec> {
let mut seen = std::collections::HashSet::new();
let mut out = Vec::new();
for source in &self.sources {
for spec in source.specs() {
if seen.insert(spec.name.clone()) {
out.push(spec);
}
}
}
if self.core.is_empty() {
return out;
}
let rank = |name: &str| self.core.iter().position(|c| c == name);
out.sort_by(|a, b| match (rank(&a.name), rank(&b.name)) {
(Some(x), Some(y)) => x.cmp(&y),
(Some(_), None) => std::cmp::Ordering::Less,
(None, Some(_)) => std::cmp::Ordering::Greater,
(None, None) => std::cmp::Ordering::Equal,
});
out
}
fn needs_approval(&self, name: &str) -> bool {
self.owner_of(name)
.is_some_and(|source| source.needs_approval(name))
}
fn pre_dispatch(&self, name: &str, args_json: &str) -> ToolDecision {
self.owner_of(name).map_or(ToolDecision::Allow, |source| {
source.pre_dispatch(name, args_json)
})
}
fn post_dispatch(&self, name: &str, args_json: &str, result_json: &str) -> Option<String> {
self.owner_of(name)
.and_then(|source| source.post_dispatch(name, args_json, result_json))
}
fn recover_unadvertised(&self, name: &str, args_json: &str) -> Vec<ToolSpec> {
let mut seen = std::collections::HashSet::new();
let mut out = Vec::new();
for source in &self.sources {
for spec in source.recover_unadvertised(name, args_json) {
if seen.insert(spec.name.clone()) {
out.push(spec);
}
}
}
out
}
fn sandbox_would_deny(&self, name: &str, args_json: &str) -> bool {
self.owner_of(name)
.is_some_and(|source| source.sandbox_would_deny(name, args_json))
}
fn required_capabilities(&self, name: &str) -> polyc_capability::CapabilitySet {
self.owner_of(name)
.map_or_else(polyc_capability::CapabilitySet::all, |source| {
source
.required_capabilities(name)
.union(crate::capability::floor_requirements(name))
})
}
fn for_worker(
&self,
scope: &polyc_agent::delegate::WorkerScope<'_>,
) -> Option<
std::result::Result<
polyc_agent::delegate::WorkerHandoff,
polyc_agent::delegate::ShareInError,
>,
> {
let rerooted: Vec<Option<polyc_agent::delegate::WorkerHandoff>> = match self
.sources
.iter()
.map(|source| source.for_worker(scope).transpose())
.collect::<std::result::Result<Vec<_>, _>>()
{
Ok(rerooted) => rerooted,
Err(err) => return Some(Err(err)),
};
if rerooted.iter().all(Option::is_none) {
return None;
}
let mut seeded = Vec::new();
let sources: Vec<Arc<dyn ToolExecutor>> = rerooted
.into_iter()
.zip(self.sources.iter())
.map(|(re, original)| {
re.map_or_else(
|| Arc::clone(original),
|handoff| {
seeded.extend(handoff.seeded);
handoff.tools
},
)
})
.collect();
Some(Ok(polyc_agent::delegate::WorkerHandoff {
tools: Arc::new(Self {
sources,
core: self.core.clone(),
}),
seeded,
}))
}
async fn execute(&self, name: &str, args_json: &str) -> String {
if let Some(source) = self.owner_of(name) {
return source.execute(name, args_json).await;
}
failure_json(
ConnectorErrorKind::Application,
&format!("unknown tool: {name}"),
)
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::pedantic, clippy::nursery, missing_docs)]
use serde_json::json;
use super::*;
#[test]
fn core_admits_only_bare_builtin_names() {
let core = vec!["grep".to_owned(), "vcs__repo_list".to_owned()];
assert!(core_admits_builtin("grep", &core));
assert!(
!core_admits_builtin("vcs__repo_list", &core),
"namespaced = grant-governed"
);
assert!(!core_admits_builtin("file_read", &core), "not declared");
}
#[test]
fn builtin_admits_gates_management_names_on_explicit_allow_only() {
let no_core: Vec<String> = Vec::new();
assert!(!builtin_admits("wallet_status", None, &no_core));
assert!(!builtin_admits("invite", None, &no_core));
let allow = vec!["wallet_status".to_owned()];
assert!(builtin_admits("wallet_status", Some(&allow), &no_core));
assert!(!builtin_admits("wallet_link", Some(&allow), &no_core));
}
#[test]
fn builtin_admits_defaults_non_management_names_on() {
let no_core: Vec<String> = Vec::new();
assert!(builtin_admits("history_search", None, &no_core));
assert!(builtin_admits("shell_exec", None, &no_core));
let allow = vec!["history_search".to_owned()];
assert!(builtin_admits("history_search", Some(&allow), &no_core));
assert!(!builtin_admits("history_peek", Some(&allow), &no_core));
}
#[test]
fn builtin_admits_core_wins_over_both_policies() {
let core = vec!["wallet_status".to_owned()];
assert!(builtin_admits("wallet_status", None, &core));
let allow: Vec<String> = Vec::new();
assert!(builtin_admits("wallet_status", Some(&allow), &core));
}
#[test]
fn connector_label_charset_is_the_k8s_resource_name_set() {
for ok in ["standup", "calc-svc", "a", "svc2"] {
assert!(is_valid_connector_label(ok), "{ok:?} must be valid");
}
for bad in [
"",
"a.b",
"a_b",
"my__svc",
"MixedCase",
"spa ce",
"emoji✨",
] {
assert!(!is_valid_connector_label(bad), "{bad:?} must be rejected");
}
}
#[derive(Debug)]
struct FakeConnector {
names: Vec<String>,
}
#[async_trait]
impl ToolExecutor for FakeConnector {
fn specs(&self) -> Vec<ToolSpec> {
self.names
.iter()
.map(|n| ToolSpec::new(n, "remote tool", json!({})))
.collect()
}
fn owns(&self, name: &str) -> bool {
self.names.iter().any(|n| n == name)
}
async fn execute(&self, _name: &str, _args_json: &str) -> String {
json!({ "ok": true }).to_string()
}
}
#[test]
fn aliased_connector_tool_classifies_through_its_owner() {
use crate::ToolRegistry;
use polyc_capability::{Capability, CapabilitySet};
let local = ToolRegistry::scoped(["shell_exec".to_owned()]);
let connector = FakeConnector {
names: vec!["file_read".to_owned()],
};
let registry = CompositeRegistry::new()
.with(Arc::new(local))
.with(Arc::new(connector));
assert_eq!(
registry.required_capabilities("file_read"),
CapabilitySet::all(),
"an aliased connector tool must classify via its owner, fail closed"
);
let local = ToolRegistry::scoped(["file_read".to_owned()]);
let connector = FakeConnector {
names: vec!["file_read".to_owned()],
};
let registry = CompositeRegistry::new()
.with(Arc::new(local))
.with(Arc::new(connector));
assert_eq!(
registry.required_capabilities("file_read"),
CapabilitySet::of(Capability::LocalRead),
"a granted local coding tool keeps its local classification"
);
}
#[test]
fn capability_floor_pins_an_under_reported_fetcher() {
use polyc_capability::{Capability, CapabilitySet};
#[derive(Debug)]
struct MisclassifyingProxy;
#[async_trait]
impl ToolExecutor for MisclassifyingProxy {
fn specs(&self) -> Vec<ToolSpec> {
vec![ToolSpec::new("web_fetch", "fetch", json!({})).read_only()]
}
fn owns(&self, name: &str) -> bool {
name == "web_fetch"
}
fn required_capabilities(&self, _name: &str) -> CapabilitySet {
CapabilitySet::of(Capability::LocalRead)
}
async fn execute(&self, _name: &str, _args_json: &str) -> String {
json!({ "ok": true }).to_string()
}
}
let registry = CompositeRegistry::new().with(Arc::new(MisclassifyingProxy));
let required = registry.required_capabilities("web_fetch");
assert!(
required.contains(Capability::ArbitraryEgress),
"the floor must keep the fetcher's requirement: {:?}",
required.names()
);
assert_eq!(
registry.required_capabilities("no_such_tool"),
CapabilitySet::all()
);
let aliased = CompositeRegistry::new().with(Arc::new(FakeConnector {
names: vec!["file_read".to_owned()],
}));
assert_eq!(
aliased.required_capabilities("file_read"),
CapabilitySet::all()
);
}
#[test]
fn resource_indicator_is_canonical_per_rfc8707() {
let token = AudienceBoundToken::new("secret", "HTTPS://Conn.Example:443/mcp#frag")
.expect("valid resource URI");
assert_eq!(token.resource(), "https://conn.example/mcp");
}
#[test]
fn token_forwarded_only_to_its_bound_audience() {
let token =
AudienceBoundToken::new("secret", "https://a.example/mcp").expect("valid resource");
assert_eq!(
token
.bearer_for("https://a.example:443/mcp")
.expect("match"),
"secret"
);
let err = token
.bearer_for("https://b.example/mcp")
.expect_err("cross-audience forward must be rejected");
assert!(
matches!(err, McpClientError::AudienceMismatch { .. }),
"expected AudienceMismatch, got {err:?}"
);
}
#[test]
fn debug_redacts_the_raw_bearer() {
let token = AudienceBoundToken::new("super-secret-bearer", "https://a.example/mcp")
.expect("valid resource");
assert!(!format!("{token:?}").contains("super-secret-bearer"));
}
#[test]
fn invalid_resource_uri_is_rejected() {
assert!(matches!(
AudienceBoundToken::new("secret", "not a url"),
Err(McpClientError::InvalidResource(_))
));
}
}