use std::{collections::HashMap, fmt, path::PathBuf, sync::Arc};
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use thiserror::Error;
use tokio_util::sync::CancellationToken;
use crate::{AgentError, ApprovalGate, ApprovalRequest, ProgressSink};
#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
pub struct ToolSpec {
pub name: String,
pub description: String,
pub parameters: Value,
}
#[derive(Debug, Clone, Copy, Serialize, Deserialize, PartialEq, Eq)]
#[serde(rename_all = "snake_case")]
pub enum ToolRisk {
ReadOnly,
Filesystem,
Process,
Delegate,
Network,
}
impl ToolRisk {
pub fn as_str(self) -> &'static str {
match self {
Self::ReadOnly => "read_only",
Self::Filesystem => "filesystem",
Self::Process => "process",
Self::Delegate => "delegate",
Self::Network => "network",
}
}
pub fn parse(value: &str) -> Option<Self> {
[
Self::ReadOnly,
Self::Filesystem,
Self::Process,
Self::Delegate,
Self::Network,
]
.into_iter()
.find(|risk| risk.as_str() == value)
}
}
#[derive(Debug, Clone)]
pub struct ToolContext {
pub call_id: String,
pub workspace: PathBuf,
pub cancellation: CancellationToken,
pub progress: ProgressSink,
pub approvals: ToolApprovals,
}
impl ToolContext {
pub fn new(workspace: PathBuf, cancellation: CancellationToken) -> Self {
Self {
call_id: String::new(),
workspace,
cancellation,
progress: ProgressSink::default(),
approvals: ToolApprovals::default(),
}
}
}
#[derive(Clone, Default)]
pub struct ToolApprovals {
gate: Option<Arc<dyn ApprovalGate>>,
call_id: String,
}
impl fmt::Debug for ToolApprovals {
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
formatter
.debug_struct("ToolApprovals")
.field("enabled", &self.gate.is_some())
.field("call_id", &self.call_id)
.finish()
}
}
impl ToolApprovals {
pub fn new(gate: Arc<dyn ApprovalGate>, call_id: impl Into<String>) -> Self {
Self {
gate: Some(gate),
call_id: call_id.into(),
}
}
#[cfg(test)]
pub(crate) fn is_enabled(&self) -> bool {
self.gate.is_some()
}
pub async fn request(
&self,
name: impl Into<String>,
risk: ToolRisk,
cwd: PathBuf,
summary: impl Into<String>,
cancellation: CancellationToken,
) -> Result<bool, AgentError> {
let Some(gate) = &self.gate else {
return Ok(false);
};
gate.approve(
ApprovalRequest {
call_id: self.call_id.clone(),
name: name.into(),
risk,
cwd,
summary: summary.into(),
},
cancellation,
)
.await
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ToolFailure {
Denied,
Cancelled,
InvalidArguments,
Unavailable,
Limit,
Failed,
UnknownTool,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ToolOutput {
pub content: String,
pub failure: Option<ToolFailure>,
pub truncated: bool,
}
impl ToolOutput {
pub fn success(content: impl Into<String>) -> Self {
Self {
content: content.into(),
failure: None,
truncated: false,
}
}
pub fn failed(failure: ToolFailure, content: impl Into<String>) -> Self {
Self {
content: content.into(),
failure: Some(failure),
truncated: false,
}
}
pub fn is_error(&self) -> bool {
self.failure.is_some()
}
}
impl From<ToolError> for ToolOutput {
fn from(error: ToolError) -> Self {
Self::failed(error.kind, error.message)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Error)]
#[error("{message}")]
pub struct ToolError {
pub kind: ToolFailure,
pub message: String,
}
impl ToolError {
pub(crate) fn new(kind: ToolFailure, message: impl Into<String>) -> Self {
Self {
kind,
message: message.into(),
}
}
pub fn invalid_arguments(message: impl Into<String>) -> Self {
Self::new(ToolFailure::InvalidArguments, message)
}
pub fn unavailable(message: impl Into<String>) -> Self {
Self::new(ToolFailure::Unavailable, message)
}
pub fn limit(message: impl Into<String>) -> Self {
Self::new(ToolFailure::Limit, message)
}
pub fn cancelled(message: impl Into<String>) -> Self {
Self::new(ToolFailure::Cancelled, message)
}
pub fn failed(message: impl Into<String>) -> Self {
Self::new(ToolFailure::Failed, message)
}
}
impl From<String> for ToolError {
fn from(message: String) -> Self {
Self::failed(message)
}
}
#[async_trait]
pub trait Tool: Send + Sync {
fn spec(&self) -> ToolSpec;
fn risk(&self, arguments: &Value) -> Result<ToolRisk, ToolError>;
fn approval_summary(&self, arguments: &Value) -> Result<String, ToolError>;
async fn execute(
&self,
arguments: Value,
context: ToolContext,
) -> Result<ToolOutput, ToolError>;
}
#[derive(Default)]
pub struct ToolRegistry {
tools: HashMap<String, Arc<dyn Tool>>,
}
impl ToolRegistry {
pub fn register(&mut self, tool: Arc<dyn Tool>) -> Result<(), ToolError> {
let name = tool.spec().name;
if self.tools.contains_key(&name) {
return Err(ToolError::failed(format!("duplicate tool name: {name}")));
}
self.tools.insert(name, tool);
Ok(())
}
pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
self.tools.get(name).cloned()
}
pub fn specs(&self) -> Vec<ToolSpec> {
let mut specs: Vec<_> = self.tools.values().map(|tool| tool.spec()).collect();
specs.sort_by(|a, b| a.name.cmp(&b.name));
specs
}
}
#[cfg(test)]
mod tests;