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