Skip to main content

atman_runtime/tools/
mod.rs

1use std::sync::Arc;
2
3use crate::tool::ToolRegistry;
4
5pub mod agent_ctrl;
6pub mod anchor;
7pub mod anchor_fs;
8pub mod bash_bg;
9pub mod context;
10pub mod flow_check;
11pub mod flow_list;
12mod flow_source;
13pub mod form;
14pub mod fs;
15pub mod git;
16pub mod git_branch;
17pub mod git_history;
18pub mod git_ops;
19pub mod git_workspace;
20pub mod git_worktree;
21pub mod help;
22pub mod hunk;
23pub mod llm_call;
24pub mod llm_classify;
25pub mod llm_extract;
26pub mod llm_generate_branches;
27pub mod memory;
28pub mod memory_stubs;
29pub mod permission;
30pub mod plan;
31pub mod preview;
32pub mod session;
33pub mod sleep;
34pub mod stdlib;
35pub mod task_ops;
36pub mod term;
37pub mod test;
38pub mod tool_output;
39pub mod web;
40
41pub fn register_tier_zero(reg: &ToolRegistry) {
42    register_tier_zero_with_rules(reg, memory_stubs::RuleFetch::new());
43}
44
45pub fn register_tier_zero_with_rules(reg: &ToolRegistry, rule_fetch: memory_stubs::RuleFetch) {
46    reg.register(Arc::new(anchor::AnchorRead));
47    reg.register(Arc::new(anchor::AnchorEdit));
48    reg.register(Arc::new(anchor::AnchorWrite));
49    reg.register(Arc::new(anchor::AnchorUndo));
50    reg.register(Arc::new(fs::FsRead));
51    reg.register(Arc::new(tool_output::OutputRead));
52    reg.register(Arc::new(fs::FsList));
53    reg.register(Arc::new(fs::FsWrite));
54    reg.register(Arc::new(fs::FsEdit));
55    reg.register(Arc::new(fs::FsGrep));
56    reg.register(Arc::new(rule_fetch));
57    reg.register(Arc::new(stdlib::ShellQuote));
58    reg.register(Arc::new(stdlib::ToJsonString));
59    reg.register(Arc::new(stdlib::ComposeEmailPreview));
60    reg.register(Arc::new(stdlib::RenderPromptXml));
61    reg.register(Arc::new(stdlib::RenderPromptMarkdown));
62    reg.register(Arc::new(stdlib::RenderPromptTerse));
63    reg.register(Arc::new(stdlib::EstimateTokens));
64    reg.register(Arc::new(stdlib::FindCompactRange));
65    reg.register(Arc::new(stdlib::ReplaceMessagesRange));
66    reg.register(Arc::new(stdlib::Len));
67    reg.register(Arc::new(stdlib::Head));
68    reg.register(Arc::new(stdlib::Tail));
69    reg.register(Arc::new(stdlib::IsEmpty));
70    reg.register(Arc::new(stdlib::Concat));
71    reg.register(Arc::new(stdlib::TextConcat));
72    reg.register(Arc::new(stdlib::ExtractToolUses));
73    reg.register(Arc::new(stdlib::DispatchAll));
74    reg.register(Arc::new(context::ContextRecordAppend));
75    reg.register(Arc::new(permission::PermissionList));
76    reg.register(Arc::new(permission::PermissionGet));
77    reg.register(Arc::new(permission::PermissionGroupTool));
78    reg.register(Arc::new(permission::PermissionUngroup));
79    reg.register(Arc::new(permission::PermissionApprove));
80    reg.register(Arc::new(permission::PermissionDeny));
81    reg.register(Arc::new(permission::PermissionDefer));
82    reg.register(Arc::new(permission::PermissionBatch));
83    reg.register(Arc::new(stdlib::MessageUser));
84    reg.register(Arc::new(stdlib::MessageAssistant));
85    reg.register(Arc::new(stdlib::MessageSystem));
86    reg.register(Arc::new(stdlib::MessageTool));
87    reg.register(Arc::new(git::GitDiff));
88    reg.register(Arc::new(git_ops::GitInit));
89    reg.register(Arc::new(git_ops::GitShow));
90    reg.register(Arc::new(git_ops::GitLog));
91    reg.register(Arc::new(git_ops::GitStatus));
92    reg.register(Arc::new(test::TestRun));
93    reg.register(Arc::new(hunk::FsEdit));
94    reg.register(Arc::new(hunk::HunkApply));
95    reg.register(Arc::new(hunk::HunkReview));
96    reg.register(Arc::new(agent_ctrl::AgentSpawn));
97    reg.register(Arc::new(agent_ctrl::AgentStatus));
98    reg.register(Arc::new(agent_ctrl::AgentOutput));
99    reg.register(Arc::new(agent_ctrl::AgentKill));
100    reg.register(Arc::new(agent_ctrl::FlowInterject));
101    reg.register(Arc::new(form::FormAsk));
102    reg.register(Arc::new(session::SessionPush));
103    reg.register(Arc::new(sleep::Sleep));
104    reg.register(Arc::new(help::HelpShow));
105    reg.register(Arc::new(flow_list::FlowList));
106    reg.register(Arc::new(flow_list::FlowSearch));
107    reg.register(Arc::new(flow_list::FlowDescribe));
108    reg.register(Arc::new(flow_check::FlowCheck));
109    reg.register(Arc::new(llm_call::LlmCallTool));
110    reg.register(Arc::new(llm_classify::LlmClassifyTool));
111    reg.register(Arc::new(llm_extract::LlmExtractTool));
112    reg.register(Arc::new(llm_generate_branches::LlmGenerateBranchesTool));
113}
114
115pub fn register_git_ops(reg: &ToolRegistry) {
116    reg.register(Arc::new(git_ops::GitAdd));
117    reg.register(Arc::new(git_ops::GitCommit));
118    reg.register(Arc::new(git_ops::GitFetch));
119    reg.register(Arc::new(git_ops::GitBranch));
120    reg.register(Arc::new(git_ops::GitPush));
121    reg.register(Arc::new(git_history::GitRestore));
122    reg.register(Arc::new(git_history::GitRevert));
123    reg.register(Arc::new(git_history::GitTagList));
124    reg.register(Arc::new(git_history::GitTagCreate));
125    reg.register(Arc::new(git_branch::GitBranchList));
126    reg.register(Arc::new(git_branch::GitBranchCreate));
127    reg.register(Arc::new(git_branch::GitBranchSwitch));
128    reg.register(Arc::new(git_branch::GitBranchRename));
129    reg.register(Arc::new(git_branch::GitBranchDelete));
130    reg.register(Arc::new(git_branch::GitRemoteList));
131    reg.register(Arc::new(git_worktree::GitWorktreeAdd));
132    reg.register(Arc::new(git_worktree::GitWorktreeList));
133    reg.register(Arc::new(git_worktree::GitWorktreeRemove));
134    reg.register(Arc::new(git_worktree::GitWorktreePrune));
135    reg.register(Arc::new(git_worktree::GitWorktreeLock));
136    reg.register(Arc::new(git_worktree::GitWorktreeUnlock));
137    reg.register(Arc::new(git_workspace::GitWorkspaceCreate));
138    reg.register(Arc::new(git_workspace::GitWorkspaceList));
139    reg.register(Arc::new(git_workspace::GitWorkspaceGet));
140    reg.register(Arc::new(git_workspace::GitWorkspaceRelease));
141    reg.register(Arc::new(git_workspace::GitWorkspaceRetain));
142    reg.register(Arc::new(git_workspace::GitWorkspacePrune));
143}
144
145pub fn register_watch(reg: &ToolRegistry) {
146    reg.register(Arc::new(crate::watch::Watch));
147    reg.register(Arc::new(crate::watch::WatcherList));
148    reg.register(Arc::new(crate::watch::WatcherUnwatch));
149    reg.register(Arc::new(crate::watch::WaitForWatcher));
150    reg.register(Arc::new(crate::watch::HasPendingInjections));
151}
152
153pub fn register_bash_bg(reg: &ToolRegistry) -> Arc<bash_bg::BgRegistry> {
154    let registry = Arc::new(bash_bg::BgRegistry::new());
155    reg.register(Arc::new(bash_bg::BashSpawn));
156    reg.register(Arc::new(bash_bg::BashStatus));
157    reg.register(Arc::new(bash_bg::BashOutput));
158    reg.register(Arc::new(bash_bg::BashKill));
159    reg.register(Arc::new(bash_bg::BashList));
160    registry
161}
162
163pub fn register_bash_bg_with_task_registry(
164    reg: &ToolRegistry,
165    task_registry: crate::task_registry::TaskRegistry,
166) -> Arc<bash_bg::BgRegistry> {
167    let registry = Arc::new(bash_bg::BgRegistry::new().with_task_registry(task_registry));
168    reg.register(Arc::new(bash_bg::BashSpawn));
169    reg.register(Arc::new(bash_bg::BashStatus));
170    reg.register(Arc::new(bash_bg::BashOutput));
171    reg.register(Arc::new(bash_bg::BashKill));
172    reg.register(Arc::new(bash_bg::BashList));
173    registry
174}
175
176pub fn register_web(reg: &ToolRegistry, config: web::WebConfig) {
177    reg.register(Arc::new(web::WebFetch::new(config)));
178}
179
180pub fn register_web_search(reg: &ToolRegistry, config: &web::SearchConfig) {
181    if let Some(provider) = web::build_search_provider(config) {
182        reg.register(Arc::new(web::WebSearch::new(provider)));
183    }
184}
185
186pub fn register_terminal(reg: &ToolRegistry) -> Arc<term::TermRegistry> {
187    let registry = Arc::new(term::TermRegistry::new());
188    reg.register(Arc::new(term::TermSpawn));
189    reg.register(Arc::new(term::TermInput));
190    reg.register(Arc::new(term::TermCapture));
191    reg.register(Arc::new(term::TermFind));
192    reg.register(Arc::new(term::TermResize));
193    reg.register(Arc::new(term::TermKill));
194    reg.register(Arc::new(term::TermList));
195    registry
196}
197
198pub fn register_terminal_with_task_registry(
199    reg: &ToolRegistry,
200    task_registry: crate::task_registry::TaskRegistry,
201) -> Arc<term::TermRegistry> {
202    let registry = Arc::new(term::TermRegistry::new().with_task_registry(task_registry));
203    reg.register(Arc::new(term::TermSpawn));
204    reg.register(Arc::new(term::TermInput));
205    reg.register(Arc::new(term::TermCapture));
206    reg.register(Arc::new(term::TermFind));
207    reg.register(Arc::new(term::TermResize));
208    reg.register(Arc::new(term::TermKill));
209    reg.register(Arc::new(term::TermList));
210    registry
211}
212
213pub fn register_preview(reg: &ToolRegistry, config: preview::PreviewConfig) {
214    reg.register(Arc::new(preview::PreviewPush::new(config)));
215}
216
217pub fn register_memory(
218    reg: &ToolRegistry,
219    todo_store: Arc<crate::memory::todo::TodoStore>,
220    confession_store: Arc<crate::memory::confession::ConfessionStore>,
221    goal_store: Arc<crate::memory::goal::GoalStore>,
222    plan_store: Arc<crate::memory::plan::PlanStore>,
223) {
224    reg.register(Arc::new(memory::MemoryTodoSet {
225        store: todo_store.clone(),
226    }));
227    reg.register(Arc::new(memory::MemoryTodoDone {
228        store: todo_store.clone(),
229    }));
230    reg.register(Arc::new(memory::MemoryTodoCancel {
231        store: todo_store.clone(),
232    }));
233    reg.register(Arc::new(memory::MemoryTodoDelete {
234        store: todo_store.clone(),
235    }));
236    reg.register(Arc::new(memory::MemoryTodoList { store: todo_store }));
237    reg.register(Arc::new(memory::MemoryConfess {
238        store: confession_store.clone(),
239    }));
240    reg.register(Arc::new(memory::MemoryFetchConfessions {
241        store: confession_store,
242    }));
243    reg.register(Arc::new(memory::MemoryGoalGet {
244        store: goal_store.clone(),
245    }));
246    reg.register(Arc::new(memory::MemoryGoalSet {
247        store: goal_store.clone(),
248    }));
249    reg.register(Arc::new(memory::MemoryGoalClear { store: goal_store }));
250    reg.register(Arc::new(memory::MemoryRecentTurns));
251    reg.register(Arc::new(memory::MemoryHistorySearch));
252    reg.register(Arc::new(memory::MemoryHistoryRead));
253    reg.register(Arc::new(memory::MemoryHistoryCount));
254    reg.register(Arc::new(plan::PlanWrite {
255        store: plan_store.clone(),
256    }));
257    reg.register(Arc::new(plan::PlanRead {
258        store: plan_store.clone(),
259    }));
260    reg.register(Arc::new(plan::PlanTick { store: plan_store }));
261}
262
263pub fn register_spec_memory(reg: &ToolRegistry, spec_store: Arc<crate::memory::spec::SpecStore>) {
264    reg.register(Arc::new(memory::MemorySpecStatus {
265        store: spec_store.clone(),
266    }));
267    reg.register(Arc::new(memory::MemorySpecMaterialize {
268        store: spec_store.clone(),
269    }));
270    reg.register(Arc::new(memory::MemorySpecUpdate {
271        store: spec_store.clone(),
272    }));
273    reg.register(Arc::new(memory::MemorySpecDeviate { store: spec_store }));
274}