cersei_tools/
lib.rs

1//! cersei-tools: Tool trait, built-in tool implementations, and permission system.
2
3pub mod apply_patch;
4pub mod ask_user;
5pub mod bash;
6pub mod bash_classifier;
7pub mod code_search;
8pub mod config_tool;
9pub mod cron;
10pub mod exa_search;
11pub mod file_edit;
12pub mod file_history;
13pub mod file_read;
14pub mod file_snapshot;
15pub mod file_watcher;
16pub mod file_write;
17pub mod git_utils;
18pub mod glob_tool;
19pub mod grep_tool;
20pub mod lsp_tool;
21pub mod notebook_edit;
22pub mod permissions;
23pub mod plan_mode;
24pub mod powershell;
25pub mod remote_trigger;
26pub mod send_message;
27pub mod skill_tool;
28pub mod skills;
29pub mod sleep;
30pub mod synthetic_output;
31pub mod tasks;
32pub mod todo_write;
33#[cfg(feature = "vms")]
34pub mod vm_tools;
35pub mod tool_primitives;
36pub mod tool_search;
37pub mod web_fetch;
38pub mod web_search;
39pub mod worktree;
40
41use async_trait::async_trait;
42use cersei_mcp::McpManager;
43use cersei_types::*;
44use serde_json::Value;
45use std::collections::HashMap;
46use std::path::PathBuf;
47use std::sync::Arc;
48
49// ─── Tool trait ──────────────────────────────────────────────────────────────
50
51#[async_trait]
52pub trait Tool: Send + Sync {
53    /// Tool name (used by the model to invoke it).
54    fn name(&self) -> &str;
55
56    /// Human-readable description shown to the model.
57    fn description(&self) -> &str;
58
59    /// JSON Schema for the tool's input parameters.
60    fn input_schema(&self) -> Value;
61
62    /// Permission level required for this tool.
63    fn permission_level(&self) -> PermissionLevel {
64        PermissionLevel::None
65    }
66
67    /// Category for grouping in tool listings.
68    fn category(&self) -> ToolCategory {
69        ToolCategory::Custom
70    }
71
72    /// Execute the tool with the given JSON input.
73    async fn execute(&self, input: Value, ctx: &ToolContext) -> ToolResult;
74
75    /// Convert to a ToolDefinition for the provider.
76    fn to_definition(&self) -> ToolDefinition {
77        ToolDefinition {
78            name: self.name().to_string(),
79            description: self.description().to_string(),
80            input_schema: self.input_schema(),
81        }
82    }
83}
84
85/// Typed tool execution trait — used with `#[derive(Tool)]`.
86#[async_trait]
87pub trait ToolExecute: Send + Sync {
88    type Input: serde::de::DeserializeOwned + schemars::JsonSchema;
89
90    async fn run(&self, input: Self::Input, ctx: &ToolContext) -> ToolResult;
91}
92
93// ─── Permission levels ───────────────────────────────────────────────────────
94
95#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
96pub enum PermissionLevel {
97    None,
98    ReadOnly,
99    Write,
100    Execute,
101    Dangerous,
102    Forbidden,
103}
104
105// ─── Tool categories ─────────────────────────────────────────────────────────
106
107#[derive(Debug, Clone, Copy, PartialEq, Eq)]
108pub enum ToolCategory {
109    FileSystem,
110    Shell,
111    Web,
112    Memory,
113    Orchestration,
114    Mcp,
115    Custom,
116}
117
118// ─── Tool result ─────────────────────────────────────────────────────────────
119
120#[derive(Debug, Clone)]
121pub struct ToolResult {
122    pub content: String,
123    pub is_error: bool,
124    pub metadata: Option<Value>,
125}
126
127impl ToolResult {
128    pub fn success(content: impl Into<String>) -> Self {
129        Self {
130            content: content.into(),
131            is_error: false,
132            metadata: None,
133        }
134    }
135
136    pub fn error(content: impl Into<String>) -> Self {
137        Self {
138            content: content.into(),
139            is_error: true,
140            metadata: None,
141        }
142    }
143
144    pub fn with_metadata(mut self, meta: Value) -> Self {
145        self.metadata = Some(meta);
146        self
147    }
148}
149
150// ─── Tool context ────────────────────────────────────────────────────────────
151
152#[derive(Clone)]
153pub struct ToolContext {
154    pub working_dir: PathBuf,
155    pub session_id: String,
156    pub permissions: Arc<dyn permissions::PermissionPolicy>,
157    pub cost_tracker: Arc<CostTracker>,
158    pub mcp_manager: Option<Arc<McpManager>>,
159    pub extensions: Extensions,
160}
161
162/// Type-map for injecting custom data into the tool context.
163#[derive(Clone, Default)]
164pub struct Extensions {
165    data: Arc<dashmap::DashMap<std::any::TypeId, Arc<dyn std::any::Any + Send + Sync>>>,
166}
167
168impl Extensions {
169    pub fn insert<T: Send + Sync + 'static>(&self, val: T) {
170        self.data.insert(std::any::TypeId::of::<T>(), Arc::new(val));
171    }
172
173    pub fn get<T: Send + Sync + 'static>(&self) -> Option<Arc<T>> {
174        self.data
175            .get(&std::any::TypeId::of::<T>())
176            .and_then(|v| Arc::clone(v.value()).downcast::<T>().ok())
177    }
178}
179
180/// Tracks cumulative token usage and cost.
181pub struct CostTracker {
182    usage: parking_lot::Mutex<Usage>,
183}
184
185impl CostTracker {
186    pub fn new() -> Self {
187        Self {
188            usage: parking_lot::Mutex::new(Usage::default()),
189        }
190    }
191
192    pub fn add(&self, usage: &Usage) {
193        self.usage.lock().merge(usage);
194    }
195
196    /// Add usage with cost estimation based on model pricing.
197    pub fn add_with_model(&self, usage: &Usage, model: &str) {
198        let mut u = usage.clone();
199        if u.cost_usd.is_none() || u.cost_usd == Some(0.0) {
200            u.cost_usd = Some(estimate_cost(model, u.input_tokens, u.output_tokens));
201        }
202        self.usage.lock().merge(&u);
203    }
204
205    pub fn current(&self) -> Usage {
206        self.usage.lock().clone()
207    }
208}
209
210impl Default for CostTracker {
211    fn default() -> Self {
212        Self::new()
213    }
214}
215
216/// Estimate USD cost from token counts based on model pricing (per 1M tokens).
217pub fn estimate_cost(model: &str, input_tokens: u64, output_tokens: u64) -> f64 {
218    let (input_per_m, output_per_m) = match model {
219        m if m.contains("gpt-5.3") => (2.0, 10.0),
220        m if m.contains("gpt-5") => (2.0, 10.0),
221        m if m.contains("gpt-4o") => (2.50, 10.0),
222        m if m.contains("gpt-4-turbo") => (10.0, 30.0),
223        m if m.starts_with("o1") => (15.0, 60.0),
224        m if m.starts_with("o3") => (10.0, 40.0),
225        m if m.contains("opus") => (15.0, 75.0),
226        m if m.contains("sonnet") => (3.0, 15.0),
227        m if m.contains("haiku") => (0.25, 1.25),
228        m if m.contains("gemini-2.0-flash") => (0.075, 0.30),
229        m if m.contains("gemini") => (1.25, 5.0),
230        m if m.contains("deepseek") => (0.27, 1.10),
231        m if m.contains("mistral-large") => (2.0, 6.0),
232        m if m.contains("llama") => (0.0, 0.0), // local/free
233        _ => (2.0, 10.0),
234    };
235    (input_tokens as f64 / 1_000_000.0) * input_per_m
236        + (output_tokens as f64 / 1_000_000.0) * output_per_m
237}
238
239// ─── Shell state (persisted across Bash invocations) ─────────────────────────
240
241#[derive(Debug, Clone, Default)]
242pub struct ShellState {
243    pub cwd: Option<PathBuf>,
244    pub env_vars: HashMap<String, String>,
245}
246
247static SHELL_STATE_REGISTRY: once_cell::sync::Lazy<
248    dashmap::DashMap<String, Arc<parking_lot::Mutex<ShellState>>>,
249> = once_cell::sync::Lazy::new(dashmap::DashMap::new);
250
251pub fn session_shell_state(session_id: &str) -> Arc<parking_lot::Mutex<ShellState>> {
252    SHELL_STATE_REGISTRY
253        .entry(session_id.to_string())
254        .or_insert_with(|| Arc::new(parking_lot::Mutex::new(ShellState::default())))
255        .clone()
256}
257
258pub fn clear_session_shell_state(session_id: &str) {
259    SHELL_STATE_REGISTRY.remove(session_id);
260}
261
262// ─── Built-in tool sets ──────────────────────────────────────────────────────
263
264/// All built-in tools (35 tools).
265pub fn all() -> Vec<Box<dyn Tool>> {
266    let mut tools: Vec<Box<dyn Tool>> = Vec::new();
267    tools.extend(filesystem());
268    tools.extend(shell());
269    tools.extend(web());
270    tools.extend(planning());
271    tools.extend(scheduling());
272    tools.extend(orchestration());
273    tools.push(Box::new(ask_user::AskUserQuestionTool));
274    tools.push(Box::new(synthetic_output::SyntheticOutputTool));
275    tools.push(Box::new(config_tool::ConfigTool));
276    tools
277}
278
279/// All coding-oriented tools (filesystem + shell + web).
280pub fn coding() -> Vec<Box<dyn Tool>> {
281    let mut tools: Vec<Box<dyn Tool>> = Vec::new();
282    tools.extend(filesystem());
283    tools.extend(shell());
284    tools.extend(web());
285    tools
286}
287
288/// File system tools: Read, Write, Edit, Glob, Grep, NotebookEdit.
289pub fn filesystem() -> Vec<Box<dyn Tool>> {
290    vec![
291        Box::new(file_read::FileReadTool),
292        Box::new(file_write::FileWriteTool),
293        Box::new(file_edit::FileEditTool),
294        Box::new(apply_patch::ApplyPatchTool),
295        Box::new(glob_tool::GlobTool),
296        Box::new(grep_tool::GrepTool),
297        Box::new(code_search::CodeSearchTool::new()),
298        Box::new(notebook_edit::NotebookEditTool),
299    ]
300}
301
302/// Shell tools: Bash, PowerShell.
303pub fn shell() -> Vec<Box<dyn Tool>> {
304    vec![
305        Box::new(bash::BashTool),
306        Box::new(powershell::PowerShellTool),
307    ]
308}
309
310/// Web tools: WebFetch, WebSearch, ExaSearch.
311pub fn web() -> Vec<Box<dyn Tool>> {
312    vec![
313        Box::new(web_fetch::WebFetchTool),
314        Box::new(web_search::WebSearchTool),
315        Box::new(exa_search::ExaSearchTool),
316    ]
317}
318
319/// Planning tools: EnterPlanMode, ExitPlanMode, TodoWrite.
320pub fn planning() -> Vec<Box<dyn Tool>> {
321    vec![
322        Box::new(plan_mode::EnterPlanModeTool),
323        Box::new(plan_mode::ExitPlanModeTool),
324        Box::new(todo_write::TodoWriteTool),
325    ]
326}
327
328/// Scheduling tools: Cron (Create/List/Delete), Sleep, RemoteTrigger.
329pub fn scheduling() -> Vec<Box<dyn Tool>> {
330    vec![
331        Box::new(cron::CronCreateTool),
332        Box::new(cron::CronListTool),
333        Box::new(cron::CronDeleteTool),
334        Box::new(sleep::SleepTool),
335        Box::new(remote_trigger::RemoteTriggerTool),
336    ]
337}
338
339/// Orchestration tools: SendMessage, Tasks, Worktree.
340pub fn orchestration() -> Vec<Box<dyn Tool>> {
341    vec![
342        Box::new(send_message::SendMessageTool),
343        Box::new(tasks::TaskCreateTool),
344        Box::new(tasks::TaskGetTool),
345        Box::new(tasks::TaskUpdateTool),
346        Box::new(tasks::TaskListTool),
347        Box::new(tasks::TaskStopTool),
348        Box::new(tasks::TaskOutputTool),
349        Box::new(worktree::EnterWorktreeTool),
350        Box::new(worktree::ExitWorktreeTool),
351    ]
352}
353
354/// No tools (for pure chat agents).
355pub fn none() -> Vec<Box<dyn Tool>> {
356    vec![]
357}
cersei_tools/lib.rs

cersei_tools/
lib.rs
