|
| 1 | +use crate::types::SpecSyncConfig; |
| 2 | +use std::fs; |
| 3 | +use std::io::{BufRead, BufReader, Write}; |
| 4 | +use std::path::Path; |
| 5 | +use std::process::{Command, Stdio}; |
| 6 | +use std::sync::mpsc; |
| 7 | +use std::time::{Duration, Instant}; |
| 8 | + |
| 9 | +const MAX_FILE_CHARS: usize = 30_000; |
| 10 | +const MAX_PROMPT_CHARS: usize = 150_000; |
| 11 | +const DEFAULT_AI_COMMAND: &str = "claude -p --output-format text"; |
| 12 | +const DEFAULT_AI_TIMEOUT_SECS: u64 = 120; |
| 13 | + |
| 14 | +/// Resolve the AI command to use. Checks config, then env, then default. |
| 15 | +pub fn resolve_ai_command(config: &SpecSyncConfig) -> Result<String, String> { |
| 16 | + // 1. Config file |
| 17 | + if let Some(cmd) = &config.ai_command { |
| 18 | + return Ok(cmd.clone()); |
| 19 | + } |
| 20 | + |
| 21 | + // 2. Environment variable |
| 22 | + if let Ok(cmd) = std::env::var("SPECSYNC_AI_COMMAND") { |
| 23 | + return Ok(cmd); |
| 24 | + } |
| 25 | + |
| 26 | + // 3. Default: check if claude CLI is available |
| 27 | + let check = Command::new("sh") |
| 28 | + .args(["-c", "command -v claude"]) |
| 29 | + .stdout(Stdio::null()) |
| 30 | + .stderr(Stdio::null()) |
| 31 | + .status(); |
| 32 | + |
| 33 | + match check { |
| 34 | + Ok(status) if status.success() => Ok(DEFAULT_AI_COMMAND.to_string()), |
| 35 | + _ => Err( |
| 36 | + "No AI command found. Install the Claude CLI, or set \"aiCommand\" in specsync.json, \ |
| 37 | + or set SPECSYNC_AI_COMMAND env var.\n\n\ |
| 38 | + Examples:\n \ |
| 39 | + \"aiCommand\": \"claude -p --output-format text\" (Claude Code CLI)\n \ |
| 40 | + \"aiCommand\": \"ollama run llama3\" (local model)\n \ |
| 41 | + \"aiCommand\": \"cat > /dev/null && echo 'test'\" (any command that reads stdin, writes stdout)" |
| 42 | + .to_string(), |
| 43 | + ), |
| 44 | + } |
| 45 | +} |
| 46 | + |
| 47 | +/// Build the prompt for spec generation. |
| 48 | +fn build_prompt( |
| 49 | + module_name: &str, |
| 50 | + source_contents: &[(String, String)], |
| 51 | + required_sections: &[String], |
| 52 | +) -> String { |
| 53 | + let sections_list = required_sections |
| 54 | + .iter() |
| 55 | + .map(|s| format!("## {s}")) |
| 56 | + .collect::<Vec<_>>() |
| 57 | + .join("\n"); |
| 58 | + |
| 59 | + let files_yaml = source_contents |
| 60 | + .iter() |
| 61 | + .map(|(path, _)| format!(" - {path}")) |
| 62 | + .collect::<Vec<_>>() |
| 63 | + .join("\n"); |
| 64 | + |
| 65 | + let mut source_block = String::new(); |
| 66 | + let mut total_chars = 0; |
| 67 | + for (path, content) in source_contents { |
| 68 | + if total_chars > MAX_PROMPT_CHARS { |
| 69 | + source_block.push_str(&format!("\n--- {path} ---\n[skipped: prompt size limit]\n")); |
| 70 | + continue; |
| 71 | + } |
| 72 | + let truncated = if content.len() > MAX_FILE_CHARS { |
| 73 | + format!( |
| 74 | + "{}\n\n[... truncated at {MAX_FILE_CHARS} chars ...]", |
| 75 | + &content[..MAX_FILE_CHARS] |
| 76 | + ) |
| 77 | + } else { |
| 78 | + content.clone() |
| 79 | + }; |
| 80 | + total_chars += truncated.len(); |
| 81 | + source_block.push_str(&format!("\n--- {path} ---\n{truncated}\n")); |
| 82 | + } |
| 83 | + |
| 84 | + format!( |
| 85 | + r#"You are a technical writer generating specification documents for software modules. |
| 86 | +Output ONLY the raw markdown spec file content. Do NOT wrap it in code fences. |
| 87 | +The spec must start with `---` YAML frontmatter and include all required sections. |
| 88 | +Be concise but thorough. Infer purpose, invariants, and error cases from the code. |
| 89 | +For the Public API section, list every public/exported symbol in markdown tables with |
| 90 | +backtick-quoted names in the first column. |
| 91 | +
|
| 92 | +Generate a spec file for the module "{module_name}". |
| 93 | +
|
| 94 | +The frontmatter must be exactly: |
| 95 | +--- |
| 96 | +module: {module_name} |
| 97 | +version: 1 |
| 98 | +status: draft |
| 99 | +files: |
| 100 | +{files_yaml} |
| 101 | +db_tables: [] |
| 102 | +depends_on: [] |
| 103 | +--- |
| 104 | +
|
| 105 | +Required markdown sections (in this order): |
| 106 | +{sections_list} |
| 107 | +
|
| 108 | +Source files: |
| 109 | +{source_block} |
| 110 | +
|
| 111 | +CRITICAL rules for the `## Public API` section: |
| 112 | +- Use markdown tables with backtick-quoted symbol names in the FIRST COLUMN |
| 113 | +- ONLY document symbols that are PUBLIC/EXPORTED from this module's external interface |
| 114 | +- Do NOT document: private functions, internal constants, private helpers, submodule names, struct fields, or implementation details |
| 115 | +- In Rust: only `pub fn`, `pub struct`, `pub enum`, `pub trait`, `pub type` that are re-exported or accessible from outside the module |
| 116 | +- In TypeScript/JS: only symbols with `export` keyword |
| 117 | +- In Python: only symbols in `__all__` or top-level non-underscore names |
| 118 | +- In Go: only capitalized names |
| 119 | +- If a symbol is private/internal (e.g. `const`, `fn` without `pub`, `mod` declarations), do NOT put it in the Public API table |
| 120 | +- Use subsection headers like `### Exported Functions`, `### Exported Types` — NOT `### Constants`, `### Per-language extractors`, `### Methods`, etc. |
| 121 | +
|
| 122 | +Other guidelines: |
| 123 | +- For `## Invariants`, list rules that must always hold based on the code |
| 124 | +- For `## Behavioral Examples`, use Given/When/Then format |
| 125 | +- For `## Error Cases`, use a table of Condition | Behavior |
| 126 | +- For `## Dependencies`, list what this module consumes from other modules |
| 127 | +- For `## Change Log`, add a single entry with today's date and "Initial spec" |
| 128 | +- Be accurate — only document what the code actually does"# |
| 129 | + ) |
| 130 | +} |
| 131 | + |
| 132 | +/// Run the AI command with the given prompt, returning stdout. |
| 133 | +/// Shows a spinner while waiting, then streams stdout lines to stderr in real time. |
| 134 | +/// Times out after `timeout_secs` seconds (default 120). |
| 135 | +fn run_ai_command(ai_command: &str, prompt: &str, timeout_secs: u64) -> Result<String, String> { |
| 136 | + let mut child = Command::new("sh") |
| 137 | + .args(["-c", ai_command]) |
| 138 | + .stdin(Stdio::piped()) |
| 139 | + .stdout(Stdio::piped()) |
| 140 | + .stderr(Stdio::piped()) |
| 141 | + .spawn() |
| 142 | + .map_err(|e| format!("Failed to start AI command: {e}"))?; |
| 143 | + |
| 144 | + if let Some(mut stdin) = child.stdin.take() { |
| 145 | + stdin |
| 146 | + .write_all(prompt.as_bytes()) |
| 147 | + .map_err(|e| format!("Failed to write to AI command stdin: {e}"))?; |
| 148 | + } |
| 149 | + |
| 150 | + // Read stdout in a background thread, streaming lines to stderr for live output |
| 151 | + let stdout_pipe = child.stdout.take().ok_or("Failed to capture stdout")?; |
| 152 | + let (tx, rx) = mpsc::channel::<String>(); |
| 153 | + let reader_thread = std::thread::spawn(move || { |
| 154 | + let reader = BufReader::new(stdout_pipe); |
| 155 | + let mut captured = String::new(); |
| 156 | + for line in reader.lines() { |
| 157 | + match line { |
| 158 | + Ok(line) => { |
| 159 | + // Stream to stderr so the user sees live progress |
| 160 | + let _ = tx.send(line.clone()); |
| 161 | + captured.push_str(&line); |
| 162 | + captured.push('\n'); |
| 163 | + } |
| 164 | + Err(_) => break, |
| 165 | + } |
| 166 | + } |
| 167 | + captured |
| 168 | + }); |
| 169 | + |
| 170 | + // Read stderr in a background thread |
| 171 | + let stderr_pipe = child.stderr.take().ok_or("Failed to capture stderr")?; |
| 172 | + let stderr_thread = std::thread::spawn(move || { |
| 173 | + let reader = BufReader::new(stderr_pipe); |
| 174 | + let mut captured = String::new(); |
| 175 | + for line in reader.lines() { |
| 176 | + match line { |
| 177 | + Ok(line) => { |
| 178 | + captured.push_str(&line); |
| 179 | + captured.push('\n'); |
| 180 | + } |
| 181 | + Err(_) => break, |
| 182 | + } |
| 183 | + } |
| 184 | + captured |
| 185 | + }); |
| 186 | + |
| 187 | + let timeout = Duration::from_secs(timeout_secs); |
| 188 | + let start = Instant::now(); |
| 189 | + let mut line_count = 0; |
| 190 | + let spinner_frames = ["⠋", "⠙", "⠹", "⠸", "⠼", "⠴", "⠦", "⠧", "⠇", "⠏"]; |
| 191 | + let mut spinner_idx = 0; |
| 192 | + let mut got_first_line = false; |
| 193 | + |
| 194 | + // Poll for lines and check timeout |
| 195 | + loop { |
| 196 | + match rx.recv_timeout(Duration::from_millis(100)) { |
| 197 | + Ok(line) => { |
| 198 | + if !got_first_line { |
| 199 | + // Clear the spinner line before printing first output line |
| 200 | + eprint!("\r\x1b[2K"); |
| 201 | + got_first_line = true; |
| 202 | + } |
| 203 | + line_count += 1; |
| 204 | + // Print live to stderr with a prefix so it's visually distinct |
| 205 | + eprintln!(" │ {line}"); |
| 206 | + } |
| 207 | + Err(mpsc::RecvTimeoutError::Timeout) => { |
| 208 | + if !got_first_line { |
| 209 | + let elapsed = start.elapsed().as_secs(); |
| 210 | + let frame = spinner_frames[spinner_idx % spinner_frames.len()]; |
| 211 | + eprint!("\r\x1b[2K {frame} Waiting for AI response... ({elapsed}s)"); |
| 212 | + let _ = std::io::stderr().flush(); |
| 213 | + spinner_idx += 1; |
| 214 | + } |
| 215 | + } |
| 216 | + Err(mpsc::RecvTimeoutError::Disconnected) => { |
| 217 | + if !got_first_line { |
| 218 | + eprint!("\r\x1b[2K"); |
| 219 | + let _ = std::io::stderr().flush(); |
| 220 | + } |
| 221 | + break; |
| 222 | + } |
| 223 | + } |
| 224 | + |
| 225 | + if start.elapsed() > timeout { |
| 226 | + eprint!("\r\x1b[2K"); |
| 227 | + let _ = std::io::stderr().flush(); |
| 228 | + let _ = child.kill(); |
| 229 | + let _ = child.wait(); |
| 230 | + return Err(format!( |
| 231 | + "AI command timed out after {timeout_secs}s ({line_count} lines received). \ |
| 232 | + Set \"aiTimeout\" in specsync.json to increase the limit." |
| 233 | + )); |
| 234 | + } |
| 235 | + } |
| 236 | + |
| 237 | + // Drain any remaining lines |
| 238 | + for line in rx.try_iter() { |
| 239 | + eprintln!(" │ {line}"); |
| 240 | + } |
| 241 | + |
| 242 | + let stdout = reader_thread |
| 243 | + .join() |
| 244 | + .map_err(|_| "stdout reader thread panicked".to_string())?; |
| 245 | + |
| 246 | + let stderr_output = stderr_thread |
| 247 | + .join() |
| 248 | + .map_err(|_| "stderr reader thread panicked".to_string())?; |
| 249 | + |
| 250 | + let status = child |
| 251 | + .wait() |
| 252 | + .map_err(|e| format!("AI command failed: {e}"))?; |
| 253 | + |
| 254 | + if !status.success() { |
| 255 | + return Err(format!( |
| 256 | + "AI command exited with {}: {}", |
| 257 | + status, |
| 258 | + stderr_output.trim() |
| 259 | + )); |
| 260 | + } |
| 261 | + |
| 262 | + if stdout.trim().is_empty() { |
| 263 | + return Err("AI command returned empty output".to_string()); |
| 264 | + } |
| 265 | + |
| 266 | + Ok(stdout) |
| 267 | +} |
| 268 | + |
| 269 | +/// Generate a spec file using AI for a given module. |
| 270 | +pub fn generate_spec_with_ai( |
| 271 | + module_name: &str, |
| 272 | + source_files: &[String], |
| 273 | + root: &Path, |
| 274 | + config: &SpecSyncConfig, |
| 275 | + ai_command: &str, |
| 276 | +) -> Result<String, String> { |
| 277 | + let mut source_contents = Vec::new(); |
| 278 | + for file in source_files { |
| 279 | + let full_path = root.join(file); |
| 280 | + let rel_path = full_path |
| 281 | + .strip_prefix(root) |
| 282 | + .map(|p| p.to_string_lossy().to_string()) |
| 283 | + .unwrap_or_else(|_| file.clone()); |
| 284 | + let content = |
| 285 | + fs::read_to_string(&full_path).map_err(|e| format!("Cannot read {file}: {e}"))?; |
| 286 | + source_contents.push((rel_path, content)); |
| 287 | + } |
| 288 | + |
| 289 | + let prompt = build_prompt(module_name, &source_contents, &config.required_sections); |
| 290 | + let timeout = config.ai_timeout.unwrap_or(DEFAULT_AI_TIMEOUT_SECS); |
| 291 | + let mut spec = run_ai_command(ai_command, &prompt, timeout)?; |
| 292 | + |
| 293 | + // Strip code fences if the model wrapped the output |
| 294 | + if spec.trim_start().starts_with("```") { |
| 295 | + let trimmed = spec.trim(); |
| 296 | + // Remove opening fence (```markdown or ```) |
| 297 | + if let Some(rest) = trimmed.strip_prefix("```markdown\n").or_else(|| trimmed.strip_prefix("```md\n")).or_else(|| trimmed.strip_prefix("```\n")) { |
| 298 | + spec = rest.to_string(); |
| 299 | + } |
| 300 | + // Remove closing fence |
| 301 | + if let Some(rest) = spec.trim_end().strip_suffix("```") { |
| 302 | + spec = rest.to_string(); |
| 303 | + } |
| 304 | + } |
| 305 | + |
| 306 | + // Validate the response has frontmatter |
| 307 | + if !spec.trim_start().starts_with("---") { |
| 308 | + return Err("AI response missing YAML frontmatter delimiters".to_string()); |
| 309 | + } |
| 310 | + |
| 311 | + Ok(spec) |
| 312 | +} |
0 commit comments