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