use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use async_trait::async_trait;
use serde_json::Value;
use crate::core::checkpoint::CheckpointStore;
use crate::core::types::{ToolOutput, ToolSpec};
pub mod builtins;
pub mod optimize;
pub mod subagent;
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Risk {
ReadOnly,
Write,
Exec,
Network,
}
pub struct PermissionRequest<'a> {
pub risk: Risk,
pub tool: &'a str,
pub key: &'a str,
pub summary: &'a str,
}
#[async_trait]
pub trait Permission: Send + Sync {
async fn request(&self, req: PermissionRequest<'_>) -> bool;
}
pub struct AutoApprove;
#[async_trait]
impl Permission for AutoApprove {
async fn request(&self, _req: PermissionRequest<'_>) -> bool {
true
}
}
pub struct DenyAll;
#[async_trait]
impl Permission for DenyAll {
async fn request(&self, _req: PermissionRequest<'_>) -> bool {
false
}
}
#[derive(Clone)]
pub struct ToolCtx {
pub cwd: PathBuf,
pub permission: Arc<dyn Permission>,
pub checkpoints: Arc<Mutex<CheckpointStore>>,
}
impl ToolCtx {
pub fn new(cwd: impl Into<PathBuf>) -> Self {
ToolCtx {
cwd: cwd.into(),
permission: Arc::new(AutoApprove),
checkpoints: Arc::new(Mutex::new(CheckpointStore::new())),
}
}
pub fn with_permission(cwd: impl Into<PathBuf>, permission: Arc<dyn Permission>) -> Self {
ToolCtx {
cwd: cwd.into(),
permission,
checkpoints: Arc::new(Mutex::new(CheckpointStore::new())),
}
}
pub fn checkpoint(&self, label: &str, path: &Path) {
if let Ok(mut store) = self.checkpoints.lock() {
store.snapshot(label, std::slice::from_ref(&path.to_path_buf()));
}
}
pub fn resolve(&self, path: impl AsRef<Path>) -> PathBuf {
let p = path.as_ref();
if p.is_absolute() {
p.to_path_buf()
} else {
self.cwd.join(p)
}
}
}
#[async_trait]
pub trait Tool: Send + Sync {
fn name(&self) -> &str;
fn description(&self) -> &str;
fn schema(&self) -> Value;
fn risk(&self) -> Risk {
Risk::ReadOnly
}
async fn run(&self, input: Value, ctx: &ToolCtx) -> ToolOutput;
fn spec(&self) -> ToolSpec {
ToolSpec {
name: self.name().to_string(),
description: self.description().to_string(),
input_schema: self.schema(),
}
}
}
#[derive(Default)]
pub struct Registry {
tools: HashMap<String, Arc<dyn Tool>>,
}
impl Registry {
pub fn new() -> Self {
Registry::default()
}
pub fn register(&mut self, tool: Arc<dyn Tool>) {
self.tools.insert(tool.name().to_string(), tool);
}
pub fn get(&self, name: &str) -> Option<Arc<dyn Tool>> {
self.tools.get(name).cloned()
}
pub fn len(&self) -> usize {
self.tools.len()
}
pub fn is_empty(&self) -> bool {
self.tools.is_empty()
}
pub fn specs(&self) -> Vec<ToolSpec> {
self.tools.values().map(|t| t.spec()).collect()
}
}
#[cfg(test)]
mod tests {
use super::*;
struct Echo;
#[async_trait]
impl Tool for Echo {
fn name(&self) -> &str {
"echo"
}
fn description(&self) -> &str {
"echo the input back"
}
fn schema(&self) -> Value {
serde_json::json!({ "type": "object", "properties": { "text": { "type": "string" } } })
}
async fn run(&self, input: Value, _ctx: &ToolCtx) -> ToolOutput {
ToolOutput::ok(input["text"].as_str().unwrap_or_default())
}
}
#[tokio::test]
async fn registry_registers_and_runs() {
let mut reg = Registry::new();
reg.register(Arc::new(Echo));
assert_eq!(reg.len(), 1);
assert_eq!(reg.specs()[0].name, "echo");
let tool = reg.get("echo").unwrap();
let out = tool
.run(serde_json::json!({ "text": "hi" }), &ToolCtx::new("."))
.await;
assert_eq!(out.text, "hi");
assert!(!out.is_error);
}
}