use crate::ai::{AiResponse, ErrorAction, ValidationError};
use serde_json::Value;
pub trait ReviewStage: Send + Sync {
fn number(&self) -> u8;
fn name(&self) -> &'static str;
fn use_log_in_context(&self) -> bool {
true
}
fn validate(&mut self, response: &AiResponse) -> Result<Value, ValidationError>;
fn format_validation_feedback(&self, violation: &str) -> String {
format!(
"\n\nPrevious attempt was rejected: {}. Please correct your output format.",
violation
)
}
fn handle_recitation_error(&mut self) -> Option<ErrorAction> {
None
}
}
macro_rules! define_standard_stage {
($struct_name:ident, $num:expr, $name:expr) => {
pub struct $struct_name;
impl ReviewStage for $struct_name {
fn number(&self) -> u8 {
$num
}
fn name(&self) -> &'static str {
$name
}
fn use_log_in_context(&self) -> bool {
!((3..=6).contains(&$num))
}
fn validate(&mut self, response: &AiResponse) -> Result<Value, ValidationError> {
validate_stages_1_to_7(response)
}
fn format_validation_feedback(&self, violation: &str) -> String {
format_validation_feedback_stages_1_to_8(violation)
}
}
};
}
define_standard_stage!(Stage1, 1, "Analyze commit main goal");
define_standard_stage!(Stage2, 2, "High-level implementation verification");
define_standard_stage!(Stage3, 3, "Execution flow verification");
define_standard_stage!(Stage4, 4, "Resource management");
define_standard_stage!(Stage5, 5, "Locking and synchronization");
define_standard_stage!(Stage6, 6, "Security audit");
define_standard_stage!(Stage7, 7, "Hardware engineer's review");
pub struct Stage8;
impl ReviewStage for Stage8 {
fn number(&self) -> u8 {
8
}
fn name(&self) -> &'static str {
"Deduplication and Consolidation"
}
fn validate(&mut self, response: &AiResponse) -> Result<Value, ValidationError> {
let parsed = parse_json_response(response)?;
if let Some(c) = parsed.get("concerns") {
if !c.is_array() {
return Err(ValidationError::FormatViolation(
"output 'concerns' is not an array".to_string(),
));
}
} else {
return Err(ValidationError::FormatViolation(
"missing 'concerns' array in output".to_string(),
));
}
if let Some(c) = parsed.get("dismissed_concerns") {
if !c.is_array() {
return Err(ValidationError::FormatViolation(
"output 'dismissed_concerns' is not an array".to_string(),
));
}
} else {
return Err(ValidationError::FormatViolation(
"missing 'dismissed_concerns' array in output".to_string(),
));
}
Ok(parsed)
}
fn format_validation_feedback(&self, violation: &str) -> String {
format_validation_feedback_stages_1_to_8(violation)
}
}
pub struct Stage9;
impl ReviewStage for Stage9 {
fn number(&self) -> u8 {
9
}
fn name(&self) -> &'static str {
"Concern/dismissed-concern conflict resolution"
}
fn validate(&mut self, response: &AiResponse) -> Result<Value, ValidationError> {
let parsed = parse_json_response(response)?;
if let Some(c) = parsed.get("concerns") {
if !c.is_array() {
return Err(ValidationError::FormatViolation(
"output 'concerns' is not an array".to_string(),
));
}
} else {
return Err(ValidationError::FormatViolation(
"missing 'concerns' array in output".to_string(),
));
}
Ok(parsed)
}
fn format_validation_feedback(&self, violation: &str) -> String {
format!(
"\n\nPrevious attempt was rejected: {}. You MUST return ONLY a JSON object containing 'concerns' array.",
violation
)
}
}
pub struct Stage10;
impl ReviewStage for Stage10 {
fn number(&self) -> u8 {
10
}
fn name(&self) -> &'static str {
"Verification and severity estimation"
}
fn validate(&mut self, response: &AiResponse) -> Result<Value, ValidationError> {
let parsed = parse_json_response(response)?;
if let Some(f) = parsed.get("findings") {
if !f.is_array() {
return Err(ValidationError::FormatViolation(
"output 'findings' is not an array".to_string(),
));
}
} else {
return Err(ValidationError::FormatViolation(
"missing 'findings' array in output".to_string(),
));
}
Ok(parsed)
}
fn format_validation_feedback(&self, violation: &str) -> String {
format!(
"\n\nPrevious attempt was rejected: {}. You MUST return ONLY a JSON object containing 'findings' array.",
violation
)
}
}
#[derive(Default)]
pub struct Stage11 {
free_form_mode: bool,
}
impl ReviewStage for Stage11 {
fn number(&self) -> u8 {
11
}
fn name(&self) -> &'static str {
"LKML-friendly report generation"
}
fn validate(&mut self, response: &AiResponse) -> Result<Value, ValidationError> {
let text = response.content.as_deref().unwrap_or("");
if self.free_form_mode {
Ok(Value::String(text.to_string()))
} else {
match validate_inline_format(text) {
Ok(_) => Ok(Value::String(text.to_string())),
Err(violation) => Err(ValidationError::FormatViolation(violation)),
}
}
}
fn format_validation_feedback(&self, violation: &str) -> String {
format!(
"Previous attempt was rejected: {}. Please correct your output format.",
violation
)
}
fn handle_recitation_error(&mut self) -> Option<ErrorAction> {
if !self.free_form_mode {
self.free_form_mode = true;
let fallback_reminder = "\n\nCRITICAL: The previous attempt failed due to a RECITATION policy violation. Do NOT quote the original patch code at all. Instead, provide a free-form summary of the findings. Start your report with a note explaining that the format is altered due to recitation restrictions. Do not use the inline quoting style `>`.";
Some(ErrorAction::RetryWithFeedback(
fallback_reminder.to_string(),
))
} else {
None
}
}
}
pub fn create_stage(stage: u8) -> Box<dyn ReviewStage> {
match stage {
1 => Box::new(Stage1),
2 => Box::new(Stage2),
3 => Box::new(Stage3),
4 => Box::new(Stage4),
5 => Box::new(Stage5),
6 => Box::new(Stage6),
7 => Box::new(Stage7),
8 => Box::new(Stage8),
9 => Box::new(Stage9),
10 => Box::new(Stage10),
11 => Box::new(Stage11::default()),
_ => panic!("Unsupported stage: {}", stage),
}
}
fn validate_stages_1_to_7(response: &AiResponse) -> Result<Value, ValidationError> {
let parsed = parse_json_response(response)?;
match required_stage_arrays(&parsed) {
Ok(_) => Ok(parsed),
Err(violation) => Err(ValidationError::FormatViolation(violation)),
}
}
fn format_validation_feedback_stages_1_to_8(violation: &str) -> String {
format!(
"\n\nPrevious attempt was rejected: {}. You MUST return ONLY a JSON object containing 'concerns' and 'dismissed_concerns' arrays. If there are no concerns and no dismissed concerns, return `{{\"concerns\": [], \"dismissed_concerns\": []}}`.",
violation
)
}
fn validate_inline_format(content: &str) -> std::result::Result<(), String> {
if content.lines().any(|l| l.trim_start().starts_with("```")) {
return Err("The output contains Markdown code blocks ('```'). It must be plain text as per `inline-template.md`.".to_string());
}
if !content.lines().any(|l| l.trim_start().starts_with(">")) {
return Err("The output does not appear to quote any code or context using '>'. Please follow the quoting style in `inline-template.md`.".to_string());
}
let has_commit_header = content
.lines()
.take(20)
.any(|l| l.trim_start().to_lowercase().starts_with("commit "));
if !has_commit_header {
return Err("The output is missing the 'commit <hash>' header. Please start with the commit details (Commit, Author, Subject) as per `inline-template.md`.".to_string());
}
let has_author_header = content
.lines()
.take(20)
.any(|l| l.trim_start().to_lowercase().starts_with("author:"));
if !has_author_header {
return Err("The output is missing the 'Author: <name>' header. Please start with the commit details (Commit, Author, Subject) as per `inline-template.md`.".to_string());
}
let has_comments = content.lines().any(|l| {
let trimmed = l.trim();
if trimmed.is_empty() || trimmed.starts_with(">") {
return false;
}
let lower = trimmed.to_lowercase();
!lower.starts_with("commit ")
&& !lower.starts_with("author:")
&& !lower.starts_with("date:")
&& !lower.starts_with("link:")
});
if !has_comments {
return Err("The output appears to lack any comments or summary. You must include a summary and interspersed comments explaining the findings.".to_string());
}
Ok(())
}
fn required_stage_arrays(value: &Value) -> std::result::Result<(&[Value], &[Value]), String> {
let concerns = value
.get("concerns")
.and_then(Value::as_array)
.ok_or_else(|| "JSON output is missing the required 'concerns' array".to_string())?;
let dismissed_concerns = value
.get("dismissed_concerns")
.and_then(Value::as_array)
.ok_or_else(|| {
"JSON output is missing the required 'dismissed_concerns' array".to_string()
})?;
Ok((concerns.as_slice(), dismissed_concerns.as_slice()))
}
fn parse_json_response(response: &AiResponse) -> Result<serde_json::Value, ValidationError> {
let raw_text = response.content.as_deref().unwrap_or("");
let cleaned = crate::utils::clean_json_string(raw_text);
let parsed: serde_json::Value = serde_json::from_str(&cleaned).unwrap_or_else(|_| {
let cands = find_json_candidates(raw_text);
cands.into_iter().last().unwrap_or(serde_json::json!({}))
});
Ok(parsed)
}
fn find_json_candidates(text: &str) -> Vec<Value> {
let mut candidates = Vec::new();
let chars: Vec<char> = text.chars().collect();
let mut i = 0;
while i < chars.len() {
if chars[i] == '{'
&& let Some(end) = find_matching_brace(&chars, i)
{
let candidate: String = chars[i..=end].iter().collect();
let clean_candidate = crate::utils::clean_json_string(&candidate);
if let Ok(v) =
serde_json::from_str(&clean_candidate).or_else(|_| serde_json::from_str(&candidate))
{
candidates.push(v);
i = end + 1;
continue;
}
}
i += 1;
}
candidates
}
fn find_matching_brace(chars: &[char], start: usize) -> Option<usize> {
let mut depth = 0;
let mut in_string = false;
let mut escape = false;
for (i, c) in chars.iter().enumerate().skip(start) {
if in_string {
if escape {
escape = false;
} else if *c == '\\' {
escape = true;
} else if *c == '"' {
in_string = false;
}
} else if *c == '"' {
in_string = true;
} else if *c == '{' {
depth += 1;
} else if *c == '}' {
depth -= 1;
if depth == 0 {
return Some(i);
}
}
}
None
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn test_required_stage_arrays_accepts_empty_arrays() {
let output = json!({
"concerns": [],
"dismissed_concerns": []
});
let (concerns, dismissed_concerns) = required_stage_arrays(&output).unwrap();
assert!(concerns.is_empty());
assert!(dismissed_concerns.is_empty());
}
#[test]
fn test_required_stage_arrays_rejects_missing_dismissed_concerns() {
let output = json!({
"concerns": []
});
let err = required_stage_arrays(&output).unwrap_err();
assert!(err.contains("'dismissed_concerns'"));
}
}