From 506c8ef03c939d21adb5c67a2b5959a46a6915b7 Mon Sep 17 00:00:00 2001 From: ajianaz Date: Thu, 27 Aug 2026 15:52:23 +0700 Subject: [PATCH] feat(review): inject enclosing control-flow scope into review prompts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Review run 2 on uteke#1084 flagged a false positive: the handler validated author_type up front and set_author_type sat in a SHARED match arm fed by two producers, but build_review_prompt shipped only raw diff hunks — the LLM reconstructed branch structure from diff shape and got it wrong (#523). - new engine::enclosing module: for hunks whose add/context lines touch branching constructs, extract the enclosing function from the post-image file (brace-balance heuristic), clamped to 120 lines with head+tail windowing so shared arms stay visible without token blowup - gated injection: new files, deletions, binaries, non-branching hunks, and unreadable files are skipped - always-on prompt guardrail forbidding reachability claims unless verified against surrounding code (stage 2 of the issue proposal) - stage 3 (call-graph reachability cross-check) intentionally left out Regression tests: acceptance fixture mirrors the shared-arm case (both producers visible in injected context); negative case without branching; new-file skip; clamp bound; guardrail presence. Signed-off-by: ajianaz --- src/engine/enclosing.rs | 330 ++++++++++++++++++++++++++++++++++++++++ src/engine/llm.rs | 53 ++++++- src/engine/mod.rs | 1 + 3 files changed, 376 insertions(+), 8 deletions(-) create mode 100644 src/engine/enclosing.rs diff --git a/src/engine/enclosing.rs b/src/engine/enclosing.rs new file mode 100644 index 0000000..c6a11a2 --- /dev/null +++ b/src/engine/enclosing.rs @@ -0,0 +1,330 @@ +//! Enclosing-scope context for review prompts (#523). +//! +//! Diff hunks alone cannot answer control-flow claims ("branch X never +//! reaches call Y") — the enclosing `match` arm, its producers, and the +//! preceding if/else usually sit outside the hunk. This module extracts the +//! enclosing function/block from the POST-IMAGE file on disk for hunks that +//! touch branching constructs, bounded to avoid token blowup. + +use std::sync::LazyLock; + +use crate::engine::diff_parser::{DiffLineType, parse_diff}; + +/// Max lines of surrounding code emitted per hunk's enclosing block. +pub const MAX_CONTEXT_LINES: usize = 120; + +/// Hunks whose ADDED lines touch any of these get enclosing-scope context. +static RE_BRANCHING: LazyLock = LazyLock::new(|| { + regex::Regex::new(r"(\bmatch\b|\bif\b|\belse\b|=>|\breturn\b|\.\?|\?\s*;)").unwrap() +}); + +/// A labeled post-image snippet for one hunk. +#[derive(Debug, Clone)] +pub struct EnclosingSnippet { + /// File the snippet was read from (post-image path). + pub file: String, + /// 1-based inclusive start line in the post-image file. + pub start_line: usize, + /// Lines actually emitted. + pub lines: usize, +} + +/// Extract enclosing-scope snippets for a unified diff. +/// +/// Only files with a usable post-image on disk are considered (paths are +/// relative to `root`). New files are skipped — their whole content is +/// already inside the diff. One snippet per qualifying hunk, each capped at +/// [`MAX_CONTEXT_LINES`]. +pub fn extract_enclosing_snippets(diff: &str, root: &std::path::Path) -> Vec { + let mut out = Vec::new(); + for chunk in parse_diff(diff) { + if chunk.is_binary || chunk.is_deleted || chunk.is_new { + continue; + } + let Some(path) = &chunk.new_path else { + continue; + }; + + for hunk in &chunk.chunks { + // Gate: only spend tokens when branching is involved (#523). + // A call ADDED inside a shared arm carries the branch structure + // in its CONTEXT lines (e.g. `Ok(id) => {`), so scan both. + let touches_branching = hunk + .lines + .iter() + .any(|l| l.line_type != DiffLineType::Remove && RE_BRANCHING.is_match(&l.content)); + if !touches_branching { + continue; + } + + let Ok(content) = std::fs::read_to_string(root.join(path)) else { + continue; + }; + let lines: Vec<&str> = content.lines().collect(); + + // First post-image line inside this hunk. + let hit = hunk + .lines + .iter() + .filter(|l| l.line_type != DiffLineType::Remove) + .find_map(|l| l.new_line_no) + .unwrap_or(hunk.new_start); + + if let Some((start_idx, end_idx)) = enclosing_block_span(&lines, hit as usize) { + out.push(EnclosingSnippet { + file: path.clone(), + start_line: start_idx + 1, + lines: end_idx - start_idx + 1, + }); + } + } + } + out +} + +/// Render snippets as a prompt section with their post-image code attached. +/// +/// Takes a resolver so this module never does I/O twice — the caller reads +/// each snippet's file once and hands over its lines. +pub fn render_for_prompt( + snippets: &[EnclosingSnippet], + file_lines: impl Fn(&str) -> Option>, +) -> String { + if snippets.is_empty() { + return String::new(); + } + let mut s = String::from( + "Surrounding code from the post-image — reference only, NOT part of the diff; \ + verify branch structure here before making reachability claims:\n", + ); + for sn in snippets { + let Some(all) = file_lines(&sn.file) else { + continue; + }; + let start = sn.start_line.saturating_sub(1); + let end = (start + sn.lines).min(all.len()); + s.push_str(&format!( + "=== {} (lines {}-{}) ===\n{}\n\n", + sn.file, + start + 1, + end, + clamp_context_text(&all[start..end].join("\n")) + )); + } + s +} + +/// Find the enclosing definition block around `hit_line` (1-based). +/// +/// Naive brace-count heuristic with two passes: +/// 1. walk forward tracking open-brace line numbers; prefer the innermost +/// opener whose line looks like a definition signature (`fn`, `def`, +/// `func`, `function`); fall back to the innermost opener. +/// 2. rescan from that opener to its matching close brace. +/// +/// Braces inside string literals may skew counts — acceptable for a +/// best-effort context gate, never used for correctness decisions. +fn enclosing_block_span(lines: &[&str], hit_line: usize) -> Option<(usize, usize)> { + if lines.is_empty() || hit_line == 0 || hit_line > lines.len() { + return None; + } + + const DEF_HINTS: [&str; 4] = ["fn ", "func ", "def ", "function"]; + + let mut opens: Vec<(usize, bool)> = Vec::new(); + for (i, line) in lines.iter().enumerate() { + for c in line.chars() { + match c { + '{' => opens.push((i + 1, DEF_HINTS.iter().any(|k| line.contains(k)))), + '}' => { + opens.pop(); + } + _ => {} + } + } + if i + 1 >= hit_line { + break; + } + } + + let open_line = opens + .iter() + .rev() + .find(|(_, looks_def)| *looks_def) + .or_else(|| opens.last()) + .map(|(l, _)| *l)?; + + // Second pass: match the braces of the chosen opener. + let mut depth = 0i64; + for (i, line) in lines.iter().enumerate().skip(open_line - 1) { + for c in line.chars() { + match c { + '{' => depth += 1, + '}' => depth -= 1, + _ => {} + } + } + if depth <= 0 { + return Some((open_line - 1, i)); + } + } + Some((open_line - 1, lines.len() - 1)) +} + +/// Cap rendered content to MAX_CONTEXT_LINES, keeping the head (signature + +/// early producers) and a tail window (shared arms live at the end of long +/// functions). +pub(crate) fn clamp_context_text(text: &str) -> String { + let count = text.lines().count(); + if count <= MAX_CONTEXT_LINES { + return text.to_string(); + } + let head = MAX_CONTEXT_LINES * 2 / 3; + let tail = MAX_CONTEXT_LINES - head - 1; // minus marker line + let mut kept: Vec<&str> = text.lines().take(head).collect(); + kept.push("…"); + kept.extend(text.lines().skip(count - tail)); + kept.join("\n") +} + +#[cfg(test)] +mod tests { + use super::*; + + /// Acceptance fixture (#523): a call added inside a SHARED match arm fed + /// by two producers. The hunk alone shows only the arm; the enclosing + /// context must expose both producers so review cannot claim the call is + /// missing on one branch. + const FIXTURE: &str = r#"handlers.rs"#; + const FIXTURE_BODY: &str = r#"use crate::store::Store; + +fn remember(store: &mut Store, text: &str) -> usize { store.remember(text) } + +fn remember_with_contradiction(store: &mut Store, text: &str) -> (usize, bool) { + let id = store.remember(text); + (id, true) +} + +fn handle_request(store: &mut Store, detect_contradiction: bool) -> usize { + let author_type = "user"; + if !author_type.is_empty() { /* validated up front */ } + let result = if detect_contradiction { + remember_with_contradiction(store, "note").map(|(id, _)| id) + } else { + Ok(remember(store, "note")) + }; + match result { + Ok(id) => { + store.set_author_type(id); + id + } + Err(_) => 0, + } +} +"#; + + fn write_fixture(dir: &std::path::Path) { + std::fs::write(dir.join(FIXTURE), FIXTURE_BODY).unwrap(); + } + + fn fixture_diff() -> String { + // Adds store.set_author_type(id) inside the shared Ok(id) arm. + let body_line = |n: usize| FIXTURE_BODY.lines().nth(n - 1).unwrap_or(""); + // The added line is at post-image line 22; keep a small hunk window. + let mut d = String::from("--- a/handlers.rs\n+++ b/handlers.rs\n"); + d.push_str("@@ -19,6 +19,7 @@\n"); + for n in 19..=21 { + d.push_str(&format!(" {}\n", body_line(n))); + } + d.push_str("+ store.set_author_type(id);\n"); + for n in 22..=24 { + d.push_str(&format!(" {}\n", body_line(n))); + } + d + } + + #[test] + fn acceptance_shared_match_arm_context_included() { + let tmp = tempfile::tempdir().unwrap(); + write_fixture(tmp.path()); + let snippets = extract_enclosing_snippets(&fixture_diff(), tmp.path()); + + assert!(!snippets.is_empty(), "branching hunk must get context"); + assert_eq!(snippets[0].file, FIXTURE); + + let rendered = render_for_prompt(&snippets, |f| { + std::fs::read_to_string(tmp.path().join(f)) + .map(|c| c.lines().map(String::from).collect()) + .ok() + }); + // Both producers of `result` are visible next to the shared arm: + assert!( + rendered.contains("remember_with_contradiction"), + "producer 1 must appear in surrounding code" + ); + assert!( + rendered.contains("Ok(remember("), + "producer 2 must appear in surrounding code" + ); + assert!( + rendered.contains("set_author_type"), + "the changed arm itself" + ); + assert!(rendered.starts_with("Surrounding code")); + } + + #[test] + fn no_context_without_branching_lines() { + let tmp = tempfile::tempdir().unwrap(); + write_fixture(tmp.path()); + // Pure arithmetic change — no branching keywords in ADDED lines. + let diff = "--- a/handlers.rs\n+++ b/handlers.rs\n@@ -2,3 +2,4 @@\n use crate::store::Store;\n+let total = 1 + 2;\n fn remember"; + let snippets = extract_enclosing_snippets(diff, tmp.path()); + assert!(snippets.is_empty(), "no branching → no injection"); + } + + #[test] + fn new_files_are_skipped() { + let tmp = tempfile::tempdir().unwrap(); + let diff = "--- /dev/null\n+++ b/newthing.rs\n@@ -0,0 +1,2 @@\n+fn x(a: u8) {\n+ match a { _ => {} }\n}"; + let snippets = extract_enclosing_snippets(diff, tmp.path()); + assert!(snippets.is_empty(), "new file content is already the diff"); + } + + #[test] + fn long_functions_are_clamped() { + let big: String = std::iter::once("fn huge() {".to_string()) + .chain((0..400).map(|i| format!(" let v{i} = {i};"))) + .chain(std::iter::once("}".to_string())) + .collect::>() + .join("\n"); + let tmp = tempfile::tempdir().unwrap(); + std::fs::write(tmp.path().join("big.rs"), &big).unwrap(); + let diff = format!( + "--- a/big.rs\n+++ b/big.rs\n@@ -398,4 +398,5 @@\n{}\n+ if v399 {{}}\n}}", + big.lines().nth(397).unwrap() + ); + let snippets = extract_enclosing_snippets(&diff, tmp.path()); + assert!(!snippets.is_empty()); + let rendered = render_for_prompt(&snippets, |f| { + std::fs::read_to_string(tmp.path().join(f)) + .map(|c| c.lines().map(String::from).collect()) + .ok() + }); + let emitted = rendered.lines().count(); + assert!( + emitted <= MAX_CONTEXT_LINES + 6, + "bounded emission expected, got {emitted} lines" + ); + assert!(rendered.contains('…'), "clamp marker present"); + } + + #[test] + fn guardrail_rule_is_part_of_prompt() { + let prompt = super::super::llm::build_review_prompt("d", &[], &[], None, None); + assert!( + prompt.contains(crate::engine::llm::CONTROL_FLOW_GUARDRAIL), + "guardrail text must always ship in review prompts" + ); + } +} diff --git a/src/engine/llm.rs b/src/engine/llm.rs index 016e63f..19f2692 100644 --- a/src/engine/llm.rs +++ b/src/engine/llm.rs @@ -430,7 +430,8 @@ pub async fn review_diff( Some(create_spinner("Reviewing diff…")) }; - let user_prompt = build_review_prompt(diff, focus, rules, static_context); + let enclosing = enclosing_section(diff); + let user_prompt = build_review_prompt(diff, focus, rules, static_context, Some(&enclosing)); let system_prompt = system_prompt_override.unwrap_or(REVIEW_SYSTEM_PROMPT); @@ -505,7 +506,8 @@ pub async fn review_diff_stream( system_prompt_override: Option<&str>, static_context: Option<&str>, ) -> std::result::Result { - let user_prompt = build_review_prompt(diff, focus, rules, static_context); + let enclosing = enclosing_section(diff); + let user_prompt = build_review_prompt(diff, focus, rules, static_context, Some(&enclosing)); let system_prompt = system_prompt_override.unwrap_or(REVIEW_SYSTEM_PROMPT); @@ -792,11 +794,35 @@ pub(crate) fn extract_file_paths_from_diff(diff: &str) -> Vec { /// Build the user prompt for diff review. #[allow(clippy::format_push_string)] -fn build_review_prompt( +/// Always-on prompt guardrail (#523): stop plausible-but-wrong reachability +/// claims that come from reasoning over diff hunks alone. +pub(crate) const CONTROL_FLOW_GUARDRAIL: &str = "Control-flow guardrail: do NOT claim an execution path is unreachable or \ +that a call is missing on a branch unless the surrounding code confirms it — \ +shared match/if arms are reached by every producer feeding them."; + +/// Build the enclosing-scope prompt section for a diff (#523). +/// +/// Reads post-image files relative to CWD (diff paths are repo-rooted); +/// returns an empty string when no hunk qualifies or files are unreadable. +pub(crate) fn enclosing_section(diff: &str) -> String { + let snippets = + crate::engine::enclosing::extract_enclosing_snippets(diff, std::path::Path::new(".")); + if snippets.is_empty() { + return String::new(); + } + crate::engine::enclosing::render_for_prompt(&snippets, |f| { + std::fs::read_to_string(f) + .map(|c| c.lines().map(String::from).collect()) + .ok() + }) +} + +pub(crate) fn build_review_prompt( diff: &str, focus: &[String], rules: &[String], static_context: Option<&str>, + enclosing_context: Option<&str>, ) -> String { let mut prompt = String::new(); @@ -820,6 +846,14 @@ fn build_review_prompt( } } + // Inject enclosing-scope code for branching hunks (#523) + if let Some(ctx) = enclosing_context { + if !ctx.is_empty() { + prompt.push_str(ctx); + prompt.push('\n'); + } + } + if !focus.is_empty() { prompt.push_str(&format!("Focus areas: {}\n\n", focus.join(", "))); } @@ -832,6 +866,9 @@ fn build_review_prompt( prompt.push('\n'); } + prompt.push_str(CONTROL_FLOW_GUARDRAIL); + prompt.push_str("\n\n"); + prompt.push_str("Review the following diff:\n\n```diff\n"); prompt.push_str(diff); prompt.push_str("\n```\n"); @@ -1702,34 +1739,34 @@ mod tests { #[test] fn build_prompt_basic() { - let prompt = build_review_prompt("diff content", &[], &[], None); + let prompt = build_review_prompt("diff content", &[], &[], None, None); assert!(prompt.contains("diff content")); assert!(prompt.contains("```diff")); } #[test] fn build_prompt_with_focus() { - let prompt = build_review_prompt("d", &["security".to_string()], &[], None); + let prompt = build_review_prompt("d", &["security".to_string()], &[], None, None); assert!(prompt.contains("Focus areas: security")); } #[test] fn build_prompt_with_rules() { - let prompt = build_review_prompt("d", &[], &["no unwrap".to_string()], None); + let prompt = build_review_prompt("d", &[], &["no unwrap".to_string()], None, None); assert!(prompt.contains("no unwrap")); } #[test] fn build_prompt_contains_file_paths() { let diff = "diff --git a/src/main.rs b/src/main.rs\n--- a/src/main.rs\n+++ b/src/main.rs\n@@ -1 +1 @@\n- old\n+ new"; - let prompt = build_review_prompt(diff, &[], &[], None); + let prompt = build_review_prompt(diff, &[], &[], None, None); assert!(prompt.contains("Valid files in this diff:")); assert!(prompt.contains("src/main.rs")); } #[test] fn build_prompt_no_file_paths_for_empty_diff() { - let prompt = build_review_prompt("no diff headers here", &[], &[], None); + let prompt = build_review_prompt("no diff headers here", &[], &[], None, None); assert!(!prompt.contains("Valid files in this diff:")); } diff --git a/src/engine/mod.rs b/src/engine/mod.rs index 9f003f8..9c6d5f7 100644 --- a/src/engine/mod.rs +++ b/src/engine/mod.rs @@ -6,6 +6,7 @@ pub mod context; pub mod db_writer; pub mod debt_tracker; pub mod diff_parser; +pub mod enclosing; pub mod index_bridge; pub mod index_scanner; pub mod language_analyzer;