use std::sync::{Arc, Mutex};
use agy_bridge::{
config::CapabilitiesConfig,
hooks::{HookEntry, HookResult},
prelude::*,
};
#[tokio::main]
async fn main() -> Result<(), Box<dyn std::error::Error>> {
agy_bridge::load_dotenv();
println!("Creating Docstring Maintenance Agent...\n");
let bridge = AgyBridge::builder().build()?;
let target_dir = std::env::current_dir()?;
let target_dir_str = target_dir.to_string_lossy().to_string();
let edited_files: Arc<Mutex<Vec<String>>> = Arc::new(Mutex::new(Vec::new()));
let edited_clone = Arc::clone(&edited_files);
let mut hook_runner = Hooks::new();
hook_runner.on_pre_tool_call_decide("py_policy", move |ctx| {
if ctx.tool_name != "edit_file" {
return HookResult::allow();
}
let path = ctx
.tool_args
.as_object()
.and_then(|o| o.get("file_path").or_else(|| o.get("path")))
.and_then(serde_json::Value::as_str)
.map_or("", |s| s.strip_prefix("file://").unwrap_or(s));
let is_py = std::path::Path::new(path)
.extension()
.is_some_and(|ext| ext.eq_ignore_ascii_case("py"));
if is_py && path.starts_with(&target_dir_str) {
HookResult::allow()
} else {
HookResult::deny("Only .py files in the target directory may be edited.")
}
});
hook_runner.on_post_tool_call("edit_tracker", move |ctx| {
if ctx.tool_name == "edit_file" {
let path = ctx
.tool_args
.as_object()
.and_then(|o| o.get("file_path").or_else(|| o.get("path")))
.and_then(serde_json::Value::as_str)
.unwrap_or("<unknown>")
.to_string();
println!(" ✏️ Edited: {path}");
edited_clone.lock().unwrap().push(path);
}
});
let hook_entries = [
HookEntry {
name: "py_policy".to_string(),
point: HookPoint::PreToolCallDecide,
callback_id: "py_policy".to_string(),
},
HookEntry {
name: "edit_tracker".to_string(),
point: HookPoint::PostToolCall,
callback_id: "edit_tracker".to_string(),
},
];
let system_instructions = format!(
"You are an expert Docstring Maintenance Agent.\n\
Audit all Python files in the target directory and ensure all public \
symbols have Google-style docstrings. Add or update docstrings as needed.\n\
Do NOT modify implementation code — only docstring blocks.\n\
Target directory: {}",
target_dir.display()
);
let config = AgentConfig::builder()
.system_instructions(system_instructions)
.workspaces(std::slice::from_ref(&target_dir))
.hooks(hook_entries.clone())
.capabilities(CapabilitiesConfig::default())
.policies(vec![PolicyRule::AllowAll])
.build();
let agent = bridge.agent(config).await?;
let prompt = "Audit all Python files and ensure public symbols have Google-style docstrings.";
println!("Prompt: {prompt}\n");
let response_text = agent.chat(prompt).await?.text().await?;
println!("{response_text}");
{
let files = edited_files.lock().unwrap();
println!("\n--- Edit Summary: {} files modified ---", files.len());
for f in files.iter() {
println!(" • {f}");
}
}
agent.shutdown().await?;
Ok(())
}