use rustc_hash::FxHashSet;
use std::path::Path;
use anyhow::{Result, anyhow};
use serde_json::{Value, json};
use crate::config::constants::tools;
use crate::config::mcp::McpAllowListConfig;
use crate::tool_policy::{ToolExecutionDecision, ToolPolicy, ToolPolicyManager};
use crate::tools::names::canonical_tool_name;
use crate::tools::tool_intent::{unified_file_action_is, unified_search_action_is};
use super::ToolPermissionDecision;
use super::risk_scorer::{RiskLevel, ToolRiskContext, ToolRiskScorer, ToolSource, WorkspaceTrust};
#[derive(Clone, Default)]
pub(super) struct ToolPolicyGateway {
tool_policy: Option<ToolPolicyManager>,
preapproved_tools: FxHashSet<String>,
full_auto_allowlist: Option<FxHashSet<String>>,
enforce_safe_mode_prompts: bool,
}
impl ToolPolicyGateway {
pub async fn new(workspace_root: &Path) -> Self {
let tool_policy = match ToolPolicyManager::new_with_workspace(workspace_root).await {
Ok(manager) => Some(manager),
Err(err) => {
tracing::warn!(%err, "Failed to initialize tool policy manager");
None
}
};
Self {
tool_policy,
preapproved_tools: FxHashSet::default(),
full_auto_allowlist: None,
enforce_safe_mode_prompts: false,
}
}
pub fn with_policy_manager(manager: ToolPolicyManager) -> Self {
Self {
tool_policy: Some(manager),
preapproved_tools: FxHashSet::default(),
full_auto_allowlist: None,
enforce_safe_mode_prompts: false,
}
}
pub fn set_enforce_safe_mode_prompts(&mut self, enabled: bool) {
self.enforce_safe_mode_prompts = enabled;
}
fn requires_safe_mode_prompt(&self, safe_mode_prompt: bool) -> bool {
self.enforce_safe_mode_prompts && safe_mode_prompt
}
pub fn has_policy_manager(&self) -> bool {
self.tool_policy.is_some()
}
pub async fn sync_available_tools(&mut self, mut available: Vec<String>, mcp_keys: &[String]) {
available.extend(mcp_keys.iter().cloned());
available.sort();
available.dedup();
if let Some(ref mut policy) = self.tool_policy
&& let Err(err) = policy.update_available_tools(available).await
{
tracing::warn!(%err, "Failed to update tool policies");
}
}
pub fn apply_policy_constraints(&self, name: &str, args: &Value) -> Result<Value> {
let mut args = args.clone();
let canonical = canonical_tool_name(name);
let normalized = canonical.as_ref();
let unified_file_read =
normalized == tools::UNIFIED_FILE && unified_file_action_is(&args, "read");
let unified_search_list =
normalized == tools::UNIFIED_SEARCH && unified_search_action_is(&args, "list");
let unified_search_grep =
normalized == tools::UNIFIED_SEARCH && unified_search_action_is(&args, "grep");
if let Some(constraints) = self
.tool_policy
.as_ref()
.and_then(|tp| tp.get_constraints(normalized))
.cloned()
{
let obj = args
.as_object_mut()
.ok_or_else(|| anyhow!("Error: tool arguments must be an object"))?;
if let Some(fmt) = constraints.default_response_format {
obj.entry("response_format").or_insert(json!(fmt));
}
if let Some(allowed) = constraints.allowed_modes
&& let Some(mode) = obj.get("mode").and_then(|v| v.as_str())
&& !allowed.iter().any(|m| m == mode)
{
return Err(anyhow!(
"Mode '{}' not allowed by policy for '{}'. Allowed: {}",
mode,
normalized,
allowed.join(", ")
));
}
match normalized {
n if n == tools::UNIFIED_SEARCH && unified_search_list => {
if let Some(cap) = constraints.max_items_per_call {
let requested = obj
.get("max_items")
.and_then(|v| v.as_u64())
.unwrap_or(cap as u64) as usize;
if requested > cap {
obj.insert("max_items".to_string(), json!(cap));
obj.insert(
"_policy_note".to_string(),
json!(format!("Capped max_items to {} by policy", cap)),
);
}
}
}
n if n == tools::UNIFIED_SEARCH && unified_search_grep => {
if let Some(cap) = constraints.max_results_per_call {
let requested = obj
.get("max_results")
.and_then(|v| v.as_u64())
.unwrap_or(cap as u64) as usize;
if requested > cap {
obj.insert("max_results".to_string(), json!(cap));
obj.insert(
"_policy_note".to_string(),
json!(format!("Capped max_results to {} by policy", cap)),
);
}
}
}
n if n == tools::READ_FILE || unified_file_read => {
if let Some(cap) = constraints.max_bytes_per_read {
let requested = obj
.get("max_bytes")
.and_then(|v| v.as_u64())
.unwrap_or(cap as u64) as usize;
if requested > cap {
obj.insert("max_bytes".to_string(), json!(cap));
obj.insert(
"_policy_note".to_string(),
json!(format!("Capped max_bytes to {} by policy", cap)),
);
}
}
}
_ => {}
}
}
Ok(args)
}
pub fn policy_manager_mut(&mut self) -> Result<&mut ToolPolicyManager> {
self.tool_policy
.as_mut()
.ok_or_else(|| anyhow!("Tool policy manager not available"))
}
pub fn set_policy_manager(&mut self, manager: ToolPolicyManager) {
self.tool_policy = Some(manager);
}
pub async fn set_tool_policy(&mut self, tool_name: &str, policy: ToolPolicy) -> Result<()> {
let canonical = canonical_tool_name(tool_name);
if let Some(ref mut manager) = self.tool_policy {
manager.set_policy(canonical.as_ref(), policy).await
} else {
Err(anyhow::anyhow!("Tool policy manager not initialized"))
}
}
pub async fn add_approval_cache_key(&mut self, approval_key: &str) -> Result<()> {
if let Some(ref mut manager) = self.tool_policy {
manager.add_approval_cache_key(approval_key).await
} else {
Err(anyhow::anyhow!("Tool policy manager not initialized"))
}
}
pub async fn add_approval_cache_prefix(&mut self, prefix_entry: &str) -> Result<()> {
if let Some(ref mut manager) = self.tool_policy {
manager.add_approval_cache_prefix(prefix_entry).await
} else {
Err(anyhow::anyhow!("Tool policy manager not initialized"))
}
}
pub fn has_approval_cache_key(&self, approval_key: &str) -> bool {
self.tool_policy
.as_ref()
.is_some_and(|manager| manager.has_approval_cache_key(approval_key))
}
pub fn matching_shell_approval_prefix(
&self,
command_words: &[String],
scope_signature: &str,
) -> Option<String> {
self.tool_policy.as_ref().and_then(|manager| {
manager.matching_shell_approval_prefix(command_words, scope_signature)
})
}
pub fn get_tool_policy(&self, tool_name: &str) -> ToolPolicy {
let canonical = canonical_tool_name(tool_name);
self.tool_policy
.as_ref()
.map(|tp| tp.get_policy(canonical.as_ref()))
.unwrap_or(ToolPolicy::Allow)
}
pub async fn reset_tool_policies(&mut self) -> Result<()> {
if let Some(tp) = self.tool_policy.as_mut() {
tp.reset_all_to_prompt().await
} else {
Err(anyhow!("Tool policy manager not available"))
}
}
pub async fn allow_all_tools(&mut self) -> Result<()> {
if let Some(tp) = self.tool_policy.as_mut() {
tp.allow_all_tools().await
} else {
Err(anyhow!("Tool policy manager not available"))
}
}
pub async fn deny_all_tools(&mut self) -> Result<()> {
if let Some(tp) = self.tool_policy.as_mut() {
tp.deny_all_tools().await
} else {
Err(anyhow!("Tool policy manager not available"))
}
}
pub fn print_policy_status(&self) {
if let Some(tp) = self.tool_policy.as_ref() {
tp.print_status();
} else {
tracing::warn!("Tool policy manager not available");
}
}
pub fn enable_full_auto_mode(&mut self, allowed_tools: &[String], available_tools: &[String]) {
let mut normalized: FxHashSet<String> = FxHashSet::default();
if allowed_tools
.iter()
.any(|tool| tool.trim() == tools::WILDCARD_ALL)
{
for tool in available_tools {
let canonical = canonical_tool_name(tool);
normalized.insert(canonical.into_owned());
}
} else {
for tool in allowed_tools {
let trimmed = tool.trim();
if !trimmed.is_empty() {
let canonical = canonical_tool_name(trimmed);
normalized.insert(canonical.into_owned());
}
}
}
self.full_auto_allowlist = Some(normalized);
}
pub fn disable_full_auto_mode(&mut self) {
self.full_auto_allowlist = None;
}
pub fn current_full_auto_allowlist(&self) -> Option<Vec<String>> {
self.full_auto_allowlist.as_ref().map(|set| {
let mut items: Vec<String> = set.iter().cloned().collect();
items.sort();
items
})
}
pub fn is_allowed_in_full_auto(&self, name: &str) -> bool {
let canonical = canonical_tool_name(name);
self.full_auto_allowlist
.as_ref()
.map(|allowlist| allowlist.contains(&*canonical))
.unwrap_or(true)
}
pub fn has_full_auto_allowlist(&self) -> bool {
self.full_auto_allowlist.is_some()
}
pub async fn evaluate_tool_policy(
&mut self,
name: &str,
safe_mode_prompt: bool,
default_permission: ToolPolicy,
) -> Result<ToolPermissionDecision> {
let canonical = canonical_tool_name(name);
let normalized = canonical.as_ref();
if self.requires_safe_mode_prompt(safe_mode_prompt) {
tracing::debug!(
"Tool '{}' requires prompt in safe mode (tools_policy)",
normalized
);
return Ok(ToolPermissionDecision::Prompt);
}
if let Some(allowlist) = self.full_auto_allowlist.as_ref() {
if !allowlist.contains(normalized) {
return Ok(ToolPermissionDecision::Deny);
}
if let Some(policy_manager) = self.tool_policy.as_mut() {
match policy_manager.get_policy(normalized) {
ToolPolicy::Deny => return Ok(ToolPermissionDecision::Deny),
ToolPolicy::Allow => {
self.preapproved_tools.insert(normalized.to_string());
return Ok(ToolPermissionDecision::Allow);
}
ToolPolicy::Prompt => {
return Ok(ToolPermissionDecision::Prompt);
}
}
}
self.preapproved_tools.insert(normalized.to_string());
return Ok(ToolPermissionDecision::Allow);
}
if let Some(policy_manager) = self.tool_policy.as_mut() {
match policy_manager.get_policy(normalized) {
ToolPolicy::Allow => {
self.preapproved_tools.insert(normalized.to_string());
Ok(ToolPermissionDecision::Allow)
}
ToolPolicy::Deny => Ok(ToolPermissionDecision::Deny),
ToolPolicy::Prompt => {
if Self::should_auto_approve_by_risk(normalized)
|| ToolPolicyManager::is_auto_allow_tool(normalized)
{
policy_manager
.set_policy(normalized, ToolPolicy::Allow)
.await?;
self.preapproved_tools.insert(normalized.to_string());
Ok(ToolPermissionDecision::Allow)
} else {
Ok(ToolPermissionDecision::Prompt)
}
}
}
} else {
Ok(match default_permission {
ToolPolicy::Allow => {
self.preapproved_tools.insert(normalized.to_string());
ToolPermissionDecision::Allow
}
ToolPolicy::Deny => ToolPermissionDecision::Deny,
ToolPolicy::Prompt => ToolPermissionDecision::Prompt,
})
}
}
fn should_auto_approve_by_risk(tool_name: &str) -> bool {
let ctx = ToolRiskContext::new(
tool_name.to_string(),
ToolSource::Internal,
WorkspaceTrust::Trusted,
);
let risk = ToolRiskScorer::calculate_risk(&ctx);
matches!(risk, RiskLevel::Low)
}
pub fn take_preapproved(&mut self, name: &str) -> bool {
let canonical = canonical_tool_name(name);
let was_preapproved = self.preapproved_tools.remove(&*canonical);
tracing::trace!(
"take_preapproved: tool='{}', canonical='{}', was_preapproved={}, remaining={:?}",
name,
canonical,
was_preapproved,
self.preapproved_tools
);
was_preapproved
}
pub fn preapprove(&mut self, name: &str) {
let canonical = canonical_tool_name(name);
let canonical_owned = canonical.into_owned();
self.preapproved_tools.insert(canonical_owned.clone());
tracing::trace!(
"preapprove: tool='{}', canonical='{}', preapproved_tools={:?}",
name,
canonical_owned,
self.preapproved_tools
);
}
pub async fn should_execute_tool(&mut self, name: &str) -> Result<ToolExecutionDecision> {
let canonical = canonical_tool_name(name);
if let Some(policy_manager) = self.tool_policy.as_mut() {
policy_manager.should_execute_tool(canonical.as_ref()).await
} else {
Ok(ToolExecutionDecision::Allowed)
}
}
pub async fn update_mcp_tools(
&mut self,
mcp_tool_index: &hashbrown::HashMap<String, Vec<String>>,
) -> Result<Option<McpAllowListConfig>> {
if let Some(policy_manager) = self.tool_policy.as_mut() {
policy_manager.update_mcp_tools(mcp_tool_index).await?;
return Ok(Some(policy_manager.mcp_allowlist().clone()));
}
Ok(None)
}
pub async fn persist_mcp_tool_policy(
&mut self,
provider: &str,
tool_name: &str,
policy: ToolPolicy,
) -> Result<()> {
if let Some(manager) = self.tool_policy.as_mut() {
manager
.set_mcp_tool_policy(provider, tool_name, policy)
.await?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tool_policy::{ToolConstraints, ToolPolicyConfig};
use indexmap::IndexMap;
use serde_json::json;
async fn gateway_with_constraints(
tool_name: &str,
constraints: ToolConstraints,
) -> ToolPolicyGateway {
let temp = tempfile::tempdir().expect("temp workspace");
let config_path = temp.path().join("tool-policy.json");
let config = ToolPolicyConfig {
constraints: IndexMap::from([(tool_name.to_string(), constraints)]),
..ToolPolicyConfig::default()
};
std::fs::write(
&config_path,
serde_json::to_vec_pretty(&config).expect("policy config json"),
)
.expect("policy config file");
let manager = ToolPolicyManager::new_with_config_path(config_path)
.await
.expect("policy manager");
ToolPolicyGateway::with_policy_manager(manager)
}
#[tokio::test]
async fn apply_policy_constraints_caps_inferred_unified_search_list_calls() {
let gateway = gateway_with_constraints(
tools::UNIFIED_SEARCH,
ToolConstraints {
max_items_per_call: Some(10),
..ToolConstraints::default()
},
)
.await;
let constrained = gateway
.apply_policy_constraints(
tools::UNIFIED_SEARCH,
&json!({
"path": ".",
"pattern": "*.rs",
"max_items": 25
}),
)
.expect("constrained args");
assert_eq!(constrained["max_items"], json!(10));
assert_eq!(
constrained["_policy_note"],
json!("Capped max_items to 10 by policy")
);
}
#[tokio::test]
async fn apply_policy_constraints_caps_inferred_unified_search_grep_calls() {
let gateway = gateway_with_constraints(
tools::UNIFIED_SEARCH,
ToolConstraints {
max_results_per_call: Some(15),
..ToolConstraints::default()
},
)
.await;
let constrained = gateway
.apply_policy_constraints(
tools::UNIFIED_SEARCH,
&json!({
"path": ".",
"pattern": "ToolRegistry",
"max_results": 50
}),
)
.expect("constrained args");
assert_eq!(constrained["max_results"], json!(15));
assert_eq!(
constrained["_policy_note"],
json!("Capped max_results to 15 by policy")
);
}
}