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