use crate::factory::AgentFactory;
use crate::progress::{ProgressHandler, SilentProgress};
use anyhow::{Context, Result};
use log::debug;
use std::path::{Path, PathBuf};
pub const PLAN_TEMPLATE: &str = include_str!("../prompts/plan/1_0.md");
pub struct PlanParams {
pub provider: String,
pub goal: String,
pub output: Option<String>,
pub instructions: Option<String>,
pub system_prompt: Option<String>,
pub model: Option<String>,
pub root: Option<String>,
pub auto_approve: bool,
pub add_dirs: Vec<String>,
pub progress: Box<dyn ProgressHandler>,
}
impl Default for PlanParams {
fn default() -> Self {
Self {
provider: "claude".to_string(),
goal: String::new(),
output: None,
instructions: None,
system_prompt: None,
model: None,
root: None,
auto_approve: false,
add_dirs: Vec::new(),
progress: Box::new(SilentProgress),
}
}
}
#[derive(Debug, Clone, Default)]
pub struct PlanResult {
pub text: Option<String>,
pub written_to: Option<PathBuf>,
}
pub fn build_plan_prompt(goal: &str, instructions: Option<&str>) -> String {
let context_section = String::new();
let prompt_section = match instructions {
Some(inst) => format!("## Additional Instructions\n\n{inst}"),
None => String::new(),
};
PLAN_TEMPLATE
.replace("{GOAL}", goal)
.replace("{CONTEXT_SECTION}", &context_section)
.replace("{PROMPT}", &prompt_section)
}
pub fn resolve_output_path(output: &str) -> PathBuf {
let path = PathBuf::from(output);
if path.extension().is_some() {
path
} else {
let timestamp = chrono::Utc::now().format("%Y%m%d-%H%M%S");
path.join(format!("plan-{timestamp}.md"))
}
}
pub fn validate_output_path(path: &Path) -> Result<()> {
let home_dir = match std::env::var("ZAG_USER_HOME_DIR") {
Ok(dir) => dir,
Err(_) => return Ok(()),
};
let home = PathBuf::from(&home_dir);
let canonical_home = std::fs::canonicalize(&home).unwrap_or_else(|_| home.clone());
let check_path = if path.exists() {
path.to_path_buf()
} else {
path.parent()
.map(PathBuf::from)
.unwrap_or_else(|| PathBuf::from("."))
};
let canonical = std::fs::canonicalize(&check_path).unwrap_or_else(|_| check_path.clone());
if !canonical.starts_with(&canonical_home) {
anyhow::bail!(
"Output path '{}' is outside your home directory: {}",
path.display(),
canonical_home.display()
);
}
Ok(())
}
pub async fn run_plan(params: PlanParams) -> Result<PlanResult> {
let PlanParams {
provider,
goal,
output,
instructions,
system_prompt,
model,
root,
auto_approve,
add_dirs,
progress,
} = params;
debug!("Starting plan via {provider} for goal: {goal}");
let output_path = match output {
Some(ref out) => {
let resolved = resolve_output_path(out);
validate_output_path(&resolved)?;
Some(resolved)
}
None => None,
};
let plan_prompt = build_plan_prompt(&goal, instructions.as_deref());
progress.on_spinner_start(&format!("Initializing {provider} for planning"));
let mut agent = AgentFactory::create(
&provider,
system_prompt,
model,
root.clone(),
auto_approve,
add_dirs,
)?;
progress.on_spinner_finish();
let model_name = agent.get_model().to_string();
if output_path.is_some() {
agent.set_capture_output(true);
}
progress.on_success(&format!("Plan initialized with model {model_name}"));
let agent_output = agent.run(Some(&plan_prompt)).await?;
agent.cleanup().await?;
if let Some(path) = output_path {
let plan_text = agent_output.and_then(|o| o.result).unwrap_or_default();
if let Some(parent) = path.parent() {
std::fs::create_dir_all(parent)
.with_context(|| format!("Failed to create directory: {}", parent.display()))?;
}
std::fs::write(&path, &plan_text)
.with_context(|| format!("Failed to write plan to: {}", path.display()))?;
progress.on_success(&format!("Plan written to {}", path.display()));
Ok(PlanResult {
text: Some(plan_text),
written_to: Some(path),
})
} else {
Ok(PlanResult::default())
}
}
#[cfg(test)]
#[path = "plan_tests.rs"]
mod tests;