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