use std::path::{Path, PathBuf};
pub struct ContextManager {
project_path: PathBuf,
}
const DEFAULT_INIT_MD: &str = r#"# Project Guidelines
## Core Rules
1. Annotate path at top of each code file: `# {relative_path}`
2. Prefer modifying existing files over creating new ones
3. All file operations must stay within the project directory
4. Dangerous commands require caution
## Workflow
1. Analyze requirements first, then plan
2. Small steps, frequent testing
3. Write self-contained test functions
4. Check safety before using tools
## Code Quality
- Follow language best practices
- Keep functions reasonably short
- Add necessary docstrings
- Handle errors gracefully
------
"#;
impl ContextManager {
pub fn new(project_path: &Path) -> Self {
let cm = ContextManager {
project_path: project_path.to_path_buf(),
};
let _ = std::fs::create_dir_all(cm.context_dir());
cm
}
fn context_dir(&self) -> PathBuf {
self.project_path.join(".aacode").join("context")
}
pub fn load_init_instructions(&self) -> String {
let path = self.project_path.join("init.md");
if let Ok(s) = std::fs::read_to_string(&path) {
if !s.trim().is_empty() {
return s;
}
}
let _ = std::fs::write(&path, DEFAULT_INIT_MD);
DEFAULT_INIT_MD.to_string()
}
pub fn save_large_output(&self, content: &str, name_hint: &str) -> Option<String> {
let dir = self.context_dir();
std::fs::create_dir_all(&dir).ok()?;
let name = format!(
"{}_{}.txt",
sanitize(name_hint),
uuid::Uuid::new_v4().simple()
);
let path = dir.join(name);
std::fs::write(&path, content).ok()?;
Some(path.to_string_lossy().to_string())
}
pub fn analyze_project_structure(&self) -> String {
use std::collections::BTreeMap;
let mut ext_counts: BTreeMap<String, usize> = BTreeMap::new();
let mut top_entries: Vec<String> = Vec::new();
let exclude = [
".git",
".aacode",
"node_modules",
"target",
"__pycache__",
".venv",
];
if let Ok(rd) = std::fs::read_dir(&self.project_path) {
for e in rd.flatten() {
let name = e.file_name().to_string_lossy().to_string();
if exclude.contains(&name.as_str()) {
continue;
}
let is_dir = e.path().is_dir();
top_entries.push(if is_dir { format!("{name}/") } else { name });
}
}
top_entries.sort();
let mut stack = vec![self.project_path.clone()];
let mut visited = 0usize;
while let Some(dir) = stack.pop() {
if visited > 5000 {
break;
}
let rd = match std::fs::read_dir(&dir) {
Ok(r) => r,
Err(_) => continue,
};
for e in rd.flatten() {
let p = e.path();
let fname = e.file_name().to_string_lossy().to_string();
if exclude.contains(&fname.as_str()) {
continue;
}
if p.is_dir() {
stack.push(p);
} else {
visited += 1;
let ext = p
.extension()
.map(|s| s.to_string_lossy().to_string())
.unwrap_or_else(|| "(none)".to_string());
*ext_counts.entry(ext).or_insert(0) += 1;
}
}
}
let mut out = String::from("Project structure summary:\n\nTop-level entries:\n");
for e in top_entries.iter().take(50) {
out.push_str(&format!(" {e}\n"));
}
out.push_str("\nFile counts by extension:\n");
for (ext, n) in ext_counts.iter() {
out.push_str(&format!(" .{ext}: {n}\n"));
}
if out.len() > 1500 {
out.truncate(1500);
out.push_str("\n...(truncated)");
}
out
}
}
fn sanitize(s: &str) -> String {
s.chars()
.map(|c| {
if c.is_alphanumeric() || c == '_' {
c
} else {
'_'
}
})
.collect()
}
#[cfg(test)]
mod tests {
use super::*;
fn tmp() -> PathBuf {
let d = std::env::temp_dir().join(format!(
"aacode_ctx_{}_{}",
std::process::id(),
uuid::Uuid::new_v4().simple()
));
std::fs::create_dir_all(&d).unwrap();
d
}
#[test]
fn init_md_created_default() {
let d = tmp();
let cm = ContextManager::new(&d);
let content = cm.load_init_instructions();
assert!(content.contains("Project Guidelines"));
assert!(d.join("init.md").exists());
}
#[test]
fn init_md_existing_used() {
let d = tmp();
std::fs::write(d.join("init.md"), "# Custom\nrules").unwrap();
let cm = ContextManager::new(&d);
assert!(cm.load_init_instructions().contains("Custom"));
}
#[test]
fn large_output_archived() {
let d = tmp();
let cm = ContextManager::new(&d);
let path = cm
.save_large_output("big content", "run_shell_output")
.unwrap();
assert!(std::path::Path::new(&path).exists());
assert_eq!(std::fs::read_to_string(&path).unwrap(), "big content");
}
#[test]
fn structure_summary() {
let d = tmp();
std::fs::write(d.join("main.rs"), "fn main(){}").unwrap();
std::fs::write(d.join("lib.py"), "x=1").unwrap();
std::fs::create_dir_all(d.join("src")).unwrap();
let cm = ContextManager::new(&d);
let s = cm.analyze_project_structure();
assert!(s.contains("main.rs"));
assert!(s.contains(".rs:") || s.contains("rs:"));
}
#[test]
fn large_output_sanitizes_path_hint() {
let d = tmp();
let cm = ContextManager::new(&d);
let path = cm.save_large_output("payload", "../../etc/passwd").unwrap();
let p = std::path::Path::new(&path);
assert!(p.starts_with(cm.context_dir()));
let fname = p.file_name().unwrap().to_string_lossy();
assert!(!fname.contains('/') && !fname.contains(".."));
assert!(fname.starts_with("______etc_passwd_"));
}
#[test]
fn structure_excludes_heavy_dirs() {
let d = tmp();
std::fs::create_dir_all(d.join("node_modules")).unwrap();
std::fs::write(d.join("node_modules/evil.js"), "x").unwrap();
std::fs::create_dir_all(d.join(".git")).unwrap();
std::fs::write(d.join(".git/config"), "x").unwrap();
std::fs::create_dir_all(d.join("src")).unwrap();
std::fs::write(d.join("main.rs"), "fn main(){}").unwrap();
let cm = ContextManager::new(&d);
let s = cm.analyze_project_structure();
assert!(s.contains("main.rs"));
assert!(s.contains("src/"));
assert!(!s.contains("node_modules"));
assert!(!s.contains(".git"));
}
}