use std::borrow::Cow;
use std::env;
use std::fmt;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex, atomic::AtomicBool, mpsc};
use schemars::JsonSchema;
use serde::Deserialize;
use serde_json::Value;
use crate::context::compact::InvokedSkillsMap;
use crate::infra::hook::HookManager;
use crate::infra::skill::Skill;
use crate::llm::{FunctionObject, ToolDefinition};
use crate::message_types::AskRequest;
use crate::permission::queue::PermissionQueue;
use crate::tools::tool_names;
use super::ask::AskTool;
use super::background::{BackgroundManager, TaskOutputTool};
use super::browser::BrowserTool;
use super::compact_tool::CompactTool;
#[cfg(target_os = "macos")]
use super::computer_use::ComputerUseTool;
use super::file::{EditFileTool, GlobTool, ReadFileTool, WriteFileTool};
use super::grep::GrepTool;
use super::hook::RegisterHookTool;
use super::plan::{self, EnterPlanModeTool, ExitPlanModeTool, PlanApprovalQueue, PlanModeState};
use super::session::SessionTool;
use super::shell::ShellTool;
use super::skill::LoadSkillTool;
use super::task::{TaskManager, TaskTool};
use super::todo::{TodoManager, TodoReadTool, TodoWriteTool};
use super::web_fetch::WebFetchTool;
use super::web_search::WebSearchTool;
use super::worktree::{EnterWorktreeTool, ExitWorktreeTool, WorktreeState};
pub use crate::message_types::PlanDecision;
#[derive(Debug, Clone)]
pub struct ImageData {
pub base64: String,
pub media_type: String,
}
#[derive(Debug)]
pub struct ToolResult {
pub output: String,
pub is_error: bool,
pub images: Vec<ImageData>,
pub plan_decision: PlanDecision,
}
pub trait Tool: Send + Sync {
fn name(&self) -> &str;
fn description(&self) -> Cow<'_, str>;
fn parameters_schema(&self) -> Value;
fn execute(&self, arguments: &str, cancelled: &Arc<AtomicBool>) -> ToolResult;
fn requires_confirmation(&self) -> bool {
false
}
fn confirmation_message(&self, arguments: &str) -> String {
format!("调用工具 {} 参数: {}", self.name(), arguments)
}
fn is_available(&self) -> bool {
true
}
}
pub fn schema_to_tool_params<T: JsonSchema>() -> Value {
let root = schemars::schema_for!(T);
let mut v = serde_json::to_value(root).unwrap_or_default();
let definitions = v
.as_object()
.and_then(|o| o.get("definitions").cloned())
.and_then(|d| d.as_object().cloned());
if let Some(defs) = definitions {
inline_refs(&mut v, &defs);
}
if let Some(obj) = v.as_object_mut() {
obj.remove("$schema");
obj.remove("title");
obj.remove("definitions");
}
v
}
fn inline_refs(value: &mut Value, definitions: &serde_json::Map<String, Value>) {
match value {
Value::Object(map) => {
if let Some(ref_path) = map.get("$ref").and_then(|r| r.as_str())
&& let Some(key) = ref_path.strip_prefix("#/definitions/")
&& let Some(def) = definitions.get(key)
{
*value = def.clone();
inline_refs(value, definitions);
return;
}
for v in map.values_mut() {
inline_refs(v, definitions);
}
}
Value::Array(arr) => {
for v in arr.iter_mut() {
inline_refs(v, definitions);
}
}
_ => {}
}
}
pub fn parse_tool_args<T: for<'de> Deserialize<'de>>(arguments: &str) -> Result<T, ToolResult> {
serde_json::from_str::<T>(arguments).map_err(|e| ToolResult {
output: format!("参数解析失败: {}", e),
is_error: true,
images: vec![],
plan_decision: PlanDecision::None,
})
}
pub struct ToolRegistry {
tools: Vec<Box<dyn Tool>>,
pub todo_manager: Arc<TodoManager>,
pub plan_mode_state: Arc<PlanModeState>,
#[allow(dead_code)]
pub worktree_state: Arc<WorktreeState>,
pub permission_queue: Option<Arc<PermissionQueue>>,
pub plan_approval_queue: Option<Arc<PlanApprovalQueue>>,
}
impl fmt::Debug for ToolRegistry {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
let tool_names: Vec<&str> = self.tools.iter().map(|t| t.name()).collect();
f.debug_struct("ToolRegistry")
.field("tool_names", &tool_names)
.finish()
}
}
impl ToolRegistry {
pub fn new(
skills: Vec<Skill>,
ask_tx: mpsc::Sender<AskRequest>,
background_manager: Arc<BackgroundManager>,
task_manager: Arc<TaskManager>,
hook_manager: Arc<Mutex<HookManager>>,
invoked_skills: InvokedSkillsMap,
todos_file_path: PathBuf,
) -> Self {
let todo_manager = Arc::new(TodoManager::new_with_file_path(todos_file_path));
let plan_mode_state = Arc::new(PlanModeState::new());
let worktree_state = Arc::new(WorktreeState::new());
let plan_approval_queue = Arc::new(PlanApprovalQueue::new());
let tools: Vec<Box<dyn Tool>> = vec![
Box::new(ShellTool {
manager: Arc::clone(&background_manager),
}),
Box::new(ReadFileTool),
Box::new(WriteFileTool),
Box::new(EditFileTool),
Box::new(GlobTool),
Box::new(GrepTool),
Box::new(WebFetchTool),
Box::new(WebSearchTool),
Box::new(BrowserTool),
Box::new(AskTool {
ask_tx: ask_tx.clone(),
}),
Box::new(TaskOutputTool {
manager: Arc::clone(&background_manager),
}),
Box::new(SessionTool {
manager: Arc::clone(&background_manager),
}),
Box::new(TaskTool {
manager: Arc::clone(&task_manager),
}),
Box::new(TodoWriteTool {
manager: Arc::clone(&todo_manager),
}),
Box::new(TodoReadTool {
manager: Arc::clone(&todo_manager),
}),
Box::new(CompactTool),
Box::new(RegisterHookTool { hook_manager }),
#[cfg(target_os = "macos")]
Box::new(ComputerUseTool::new()),
Box::new(EnterPlanModeTool {
plan_state: Arc::clone(&plan_mode_state),
}),
Box::new(ExitPlanModeTool {
plan_state: Arc::clone(&plan_mode_state),
ask_tx,
plan_approval_queue: Some(Arc::clone(&plan_approval_queue)),
}),
Box::new(EnterWorktreeTool {
state: Arc::clone(&worktree_state),
}),
Box::new(ExitWorktreeTool {
state: Arc::clone(&worktree_state),
}),
];
let mut registry = Self {
todo_manager: Arc::clone(&todo_manager),
plan_mode_state: Arc::clone(&plan_mode_state),
worktree_state: Arc::clone(&worktree_state),
permission_queue: None,
plan_approval_queue: None,
tools,
};
if !skills.is_empty() {
registry.register(Box::new(LoadSkillTool {
skills,
invoked_skills,
}));
}
registry
}
pub fn register(&mut self, tool: Box<dyn Tool>) {
self.tools.push(tool);
}
pub fn get(&self, name: &str) -> Option<&dyn Tool> {
self.tools
.iter()
.find(|t| t.name() == name)
.map(|t| t.as_ref())
}
pub fn execute(&self, name: &str, arguments: &str, cancelled: &Arc<AtomicBool>) -> ToolResult {
let (is_active, plan_file_path) = self.plan_mode_state.get_state();
if is_active && !plan::is_allowed_in_plan_mode(name) {
let is_plan_file_write = (name == "Write" || name == "Edit") && {
if let Some(ref plan_path) = plan_file_path {
serde_json::from_str::<Value>(arguments)
.ok()
.and_then(|v| {
v.get("path")
.or_else(|| v.get("file_path"))
.and_then(|p| p.as_str())
.map(|p| {
let input_path = Path::new(p);
let plan_path_buf = Path::new(&plan_path);
if p == plan_path {
return true;
}
if input_path.is_relative()
&& let Ok(cwd) = env::current_dir()
{
let absolute_path = cwd.join(input_path);
if let Ok(canonical_input) = absolute_path.canonicalize()
&& let Ok(canonical_plan) = plan_path_buf.canonicalize()
{
return canonical_input == canonical_plan;
}
}
false
})
})
.unwrap_or(false)
} else {
false
}
};
if !is_plan_file_write {
return ToolResult {
output: format!(
"Tool '{}' is not available in plan mode. Only read-only tools are allowed. \
Use ExitPlanMode to exit plan mode first.",
name
),
is_error: true,
images: vec![],
plan_decision: PlanDecision::None,
};
}
}
match self.get(name) {
Some(tool) => {
if !tool.is_available() {
return ToolResult {
output: format!("Tool '{}' is currently not available.", name),
is_error: true,
images: vec![],
plan_decision: PlanDecision::None,
};
}
tool.execute(arguments, cancelled)
}
None => ToolResult {
output: format!("未知工具: {}", name),
is_error: true,
images: vec![],
plan_decision: PlanDecision::None,
},
}
}
pub fn build_tools_summary_non_deferred(
&self,
disabled: &[String],
deferred: &[String],
) -> String {
let mut md = String::new();
for t in self
.tools
.iter()
.filter(|t| !disabled.iter().any(|d| d == t.name()))
.filter(|t| t.is_available())
.filter(|t| !deferred.iter().any(|d| d == t.name()))
{
let name = t.name();
md.push_str(&format!("<{}>\n", name));
let mut desc = dedent(t.description().trim());
if name == tool_names::LOAD_TOOL {
desc.push_str(&format_deferred_suffix(deferred));
}
md.push_str(&format!("description:\n{}\n", desc));
let params = json_schema_to_xml_params(&t.parameters_schema());
if !params.is_empty() {
md.push('\n');
md.push_str(¶ms);
}
md.push_str(&format!("</{}>\n\n", name));
}
md.trim_end().to_string()
}
pub fn to_llm_tools_non_deferred(
&self,
disabled: &[String],
deferred: &[String],
) -> Vec<ToolDefinition> {
let mut tools: Vec<ToolDefinition> = self
.tools
.iter()
.filter(|t| !disabled.iter().any(|d| d == t.name()))
.filter(|t| t.is_available())
.filter(|t| !deferred.iter().any(|d| d == t.name()))
.map(|t| {
let mut desc = dedent(t.description().trim());
if t.name() == tool_names::LOAD_TOOL {
desc.push_str(&format_deferred_suffix(deferred));
}
ToolDefinition {
tool_type: "function".to_string(),
function: FunctionObject {
name: t.name().to_string(),
description: Some(desc),
parameters: Some(t.parameters_schema()),
strict: None,
},
}
})
.collect();
if let Some(load_tool) = self
.tools
.iter()
.find(|t| t.name() == tool_names::LOAD_TOOL)
&& load_tool.is_available()
&& !disabled.iter().any(|d| d == load_tool.name())
&& !tools
.iter()
.any(|t| t.function.name == tool_names::LOAD_TOOL)
{
let mut desc = dedent(load_tool.description().trim());
desc.push_str(&format_deferred_suffix(deferred));
tools.push(ToolDefinition {
tool_type: "function".to_string(),
function: FunctionObject {
name: tool_names::LOAD_TOOL.to_string(),
description: Some(desc),
parameters: Some(load_tool.parameters_schema()),
strict: None,
},
});
}
tools
}
pub fn tool_names(&self) -> Vec<&str> {
self.tools.iter().map(|t| t.name()).collect()
}
pub fn build_session_state_summary(&self) -> String {
let mut parts = Vec::new();
let (plan_active, plan_file) = self.plan_mode_state.get_state();
if plan_active {
let mut s = String::from("## Session State: PLAN MODE\n\n");
s.push_str("You are currently in **Plan Mode**. Only read-only tools are available.\n");
s.push_str(
"Write your plan to the plan file, then use ExitPlanMode for user approval.\n",
);
if let Some(ref path) = plan_file {
s.push_str(&format!("Plan file: `{}`\n", path));
}
parts.push(s);
}
if let Some(session) = self.worktree_state.get_session() {
let mut s = String::from("## Session State: WORKTREE\n\n");
s.push_str("You are in an isolated git worktree.\n");
s.push_str(&format!("Branch: `{}`\n", session.branch));
s.push_str(&format!(
"Worktree path: `{}`\n",
session.worktree_path.display()
));
s.push_str(&format!(
"Original cwd: `{}`\n",
session.original_cwd.display()
));
parts.push(s);
}
if parts.is_empty() {
return String::new();
}
parts.join("\n")
}
}
fn dedent(s: &str) -> String {
let lines: Vec<&str> = s.lines().collect();
if lines.is_empty() {
return String::new();
}
let min_indent_bytes = lines
.iter()
.filter(|line| !line.trim().is_empty())
.map(|line| {
line.chars()
.take_while(|c| c.is_whitespace() && c.is_ascii())
.map(|c| c.len_utf8())
.sum::<usize>()
})
.min()
.unwrap_or(0);
lines
.iter()
.map(|line| {
if line.trim().is_empty() || min_indent_bytes == 0 {
line.to_string()
} else {
let byte_offset = line
.char_indices()
.take_while(|(i, c)| *i < min_indent_bytes && c.is_whitespace() && c.is_ascii())
.map(|(_, c)| c.len_utf8())
.sum::<usize>();
if byte_offset >= line.len() {
String::new()
} else {
line[byte_offset..].to_string()
}
}
})
.collect::<Vec<_>>()
.join("\n")
}
fn json_schema_to_xml_params(schema: &Value) -> String {
let properties = match schema.get("properties").and_then(|p| p.as_object()) {
Some(p) => p,
None => return String::new(),
};
let required: Vec<&str> = schema
.get("required")
.and_then(|r| r.as_array())
.map(|arr| arr.iter().filter_map(|v| v.as_str()).collect())
.unwrap_or_default();
let mut md = String::from("parameters:\n");
for (name, prop) in properties {
let type_str = prop
.get("type")
.and_then(|t| t.as_str())
.unwrap_or("string");
let desc = prop
.get("description")
.and_then(|d| d.as_str())
.unwrap_or("");
let req = if required.contains(&name.as_str()) {
", required"
} else {
""
};
md.push_str(&format!("- `{}` ({}{}) — {}\n", name, type_str, req, desc));
}
md
}
fn format_deferred_suffix(deferred: &[String]) -> String {
if deferred.is_empty() {
"\n\nNo deferred tools available.".to_string()
} else {
format!("\n\nCurrently deferred tools: {}", deferred.join(", "))
}
}
#[cfg(test)]
mod tests;