use anyhow::{Context, Result};
use std::{env, fs, path::{Path, PathBuf}, process::{Command as StdCommand, Stdio}};
use tokio::process::Command;
use crate::types::Task;
const MAX_DIFF_CHARS: usize = 8000;
pub struct JudgeResult {
pub passed: bool,
pub feedback: String,
}
pub struct PeerReview {
pub score: u8,
pub feedback: String,
}
pub async fn judge_task(task: &Task, judge_agent: &str, original_prompt: &str) -> Result<JudgeResult> {
let diff = gather_diff(task)
.or_else(|| read_output(task))
.unwrap_or_else(|| "(no diff or output)".to_string());
let truncated = truncate_diff(&diff, MAX_DIFF_CHARS);
let prompt = format!(
concat!(
"You are a code review judge.\n\n",
"## Original task\n{}\n\n",
"## Output\n```\n{}\n```\n\n",
"## Instructions\n",
"Review whether the output satisfies the original task.\n",
"Your FIRST line of output MUST be exactly one of:\n",
" PASS: <brief reason>\n",
" RETRY: <what needs to be fixed>\n",
"Do NOT output anything before PASS or RETRY.",
),
original_prompt, truncated,
);
let exe = env::current_exe().context("Failed to locate aid binary")?;
let output = Command::new(exe)
.args(["run", judge_agent, &prompt, "--dir", "."])
.current_dir(task.repo_path.as_deref().unwrap_or("."))
.stdout(Stdio::piped())
.stderr(Stdio::null())
.output()
.await
.context("Judge subprocess failed")?;
if !output.status.success() {
anyhow::bail!("Judge agent exited: {}", output.status);
}
parse_judge_response(&String::from_utf8_lossy(&output.stdout))
}
pub async fn peer_review_task(task: &Task, reviewer_agent: &str, original_prompt: &str) -> Result<PeerReview> {
let diff = gather_diff(task)
.or_else(|| read_output(task))
.unwrap_or_else(|| "(no diff or output)".to_string());
let truncated = truncate_diff(&diff, MAX_DIFF_CHARS);
let prompt = format!(
concat!(
"You are a code review peer.\n\n",
"## Original task\n{}\n\n",
"## Output\n```\n{}\n```\n\n",
"## Instructions\n",
"Score the output quality from 1-10 and provide brief feedback.\n",
"Your FIRST line MUST be: SCORE: <number>/10\n",
"Then provide 1-3 lines of feedback.\n",
),
original_prompt, truncated,
);
let exe = std::env::current_exe().context("Failed to locate aid binary")?;
let output = tokio::process::Command::new(exe)
.args(["run", reviewer_agent, &prompt, "--dir", "."])
.current_dir(task.repo_path.as_deref().unwrap_or("."))
.stdout(std::process::Stdio::piped())
.stderr(std::process::Stdio::null())
.output()
.await
.context("Peer review subprocess failed")?;
if !output.status.success() {
anyhow::bail!("Peer reviewer exited: {}", output.status);
}
parse_peer_review(&String::from_utf8_lossy(&output.stdout))
}
pub(crate) fn gather_diff(task: &Task) -> Option<String> {
let dir = task.worktree_path.as_deref().or(task.repo_path.as_deref())?;
if !Path::new(dir).exists() {
return None;
}
for args in [&["diff", "--no-color", "HEAD~1..HEAD"][..], &["diff", "--no-color"]] {
let output = StdCommand::new("git").current_dir(dir).args(args).output().ok()?;
if output.status.success() {
let diff = String::from_utf8_lossy(&output.stdout).into_owned();
if !diff.trim().is_empty() {
return Some(diff);
}
}
}
None
}
pub(crate) fn read_output(task: &Task) -> Option<String> {
let output_path = task.output_path.as_deref()?;
let mut candidates = vec![PathBuf::from(output_path)];
if let Some(worktree) = task.worktree_path.as_deref() {
candidates.push(Path::new(worktree).join(output_path));
}
for candidate in candidates {
if let Ok(text) = fs::read_to_string(&candidate)
&& !text.trim().is_empty()
{
return Some(text);
}
}
None
}
fn truncate_diff(diff: &str, max_chars: usize) -> &str {
if diff.len() <= max_chars {
return diff;
}
match diff[..max_chars].rfind('\n') {
Some(pos) => &diff[..pos],
None => &diff[..max_chars],
}
}
fn parse_judge_response(text: &str) -> Result<JudgeResult> {
for line in text.lines() {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let up = trimmed.to_uppercase();
let (word, passed) = if up.starts_with("PASS") {
("PASS", true)
} else if up.starts_with("RETRY") {
("RETRY", false)
} else {
continue;
};
let feedback = trimmed[word.len()..]
.trim_start_matches(|c: char| c.is_ascii_whitespace() || c == ':' || c == '-')
.trim()
.to_string();
return Ok(JudgeResult { passed, feedback });
}
aid_warn!("[aid] Judge response contained no PASS/RETRY verdict — defaulting to PASS");
Ok(JudgeResult { passed: true, feedback: "no verdict found in judge response".to_string() })
}
fn parse_peer_review(text: &str) -> Result<PeerReview> {
for line in text.lines() {
let trimmed = line.trim();
if trimmed.is_empty() {
continue;
}
let up = trimmed.to_uppercase();
if up.starts_with("SCORE:") || up.starts_with("SCORE ") {
let rest = trimmed[6..].trim().trim_start_matches(':').trim();
if let Some(num_str) = rest.split('/').next()
&& let Ok(score) = num_str.trim().parse::<u8>()
{
let score = score.min(10);
let feedback: String = text
.lines()
.skip_while(|l| !l.trim().to_uppercase().starts_with("SCORE"))
.skip(1)
.filter(|l| !l.trim().is_empty())
.take(3)
.collect::<Vec<_>>()
.join(" ");
return Ok(PeerReview { score, feedback });
}
}
}
Ok(PeerReview { score: 5, feedback: "no score found in review".to_string() })
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn parse_first_line_pass() {
let result = parse_judge_response("PASS: looks good").unwrap();
assert!(result.passed);
assert_eq!(result.feedback, "looks good");
}
#[test]
fn parse_first_line_retry() {
let result = parse_judge_response("RETRY: missing tests").unwrap();
assert!(!result.passed);
assert_eq!(result.feedback, "missing tests");
}
#[test]
fn parse_verdict_after_prose() {
let text = "Looking at the diff, I can see changes were made.\nThe implementation looks complete.\nPASS: all requirements met";
let result = parse_judge_response(text).unwrap();
assert!(result.passed);
assert_eq!(result.feedback, "all requirements met");
}
#[test]
fn parse_retry_after_reasoning() {
let text = "The task asked for tests but none were added.\nRETRY: add unit tests for the new function";
let result = parse_judge_response(text).unwrap();
assert!(!result.passed);
assert_eq!(result.feedback, "add unit tests for the new function");
}
#[test]
fn parse_no_verdict_defaults_to_pass() {
let text = "The code looks fine and all changes are appropriate.";
let result = parse_judge_response(text).unwrap();
assert!(result.passed);
assert!(result.feedback.contains("no verdict"));
}
#[test]
fn parse_empty_response_defaults_to_pass() {
let result = parse_judge_response("").unwrap();
assert!(result.passed);
}
#[test]
fn truncate_diff_within_limit() {
let short = "abc\ndef";
assert_eq!(truncate_diff(short, 100), short);
}
#[test]
fn truncate_diff_at_newline_boundary() {
let diff = "line1\nline2\nline3\nline4";
let result = truncate_diff(diff, 13);
assert_eq!(result, "line1\nline2");
}
#[test]
fn parse_peer_review_extracts_score() {
let text = "SCORE: 8/10\nGood implementation, clean code.";
let review = parse_peer_review(text).unwrap();
assert_eq!(review.score, 8);
assert!(review.feedback.contains("Good implementation"));
}
#[test]
fn parse_peer_review_no_score_defaults_to_5() {
let review = parse_peer_review("The code looks fine.").unwrap();
assert_eq!(review.score, 5);
}
}