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