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