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