Skip to main content

differential_engine/grouping/
mod.rs

1//! The grouping stage: an LLM merges and labels shape-class ids — never hunks
2//! (ADR 0001) — behind the `LlmBackend` abstraction (ADR 0016), with a coverage
3//! audit that back-fills anything the model drops (invariant 5) and a
4//! content-hash cache that pins groupings (ADR 0009).
5//!
6//! Mechanical pieces the model never sees or cannot override:
7//! - classes living entirely in generated files are pre-assigned to the noise
8//!   tier and never reach the payload (ADR 0006);
9//! - classes touching a rename below 95% similarity can never stay in a skim
10//!   group — they are extracted into a synthesized focus group (ADR 0003).
11
12mod assemble;
13mod key;
14mod parse;
15mod payload;
16
17use std::collections::{HashMap, HashSet};
18
19use crate::llm::LlmBackend;
20use crate::ports::{ArtefactStore, GroupingCache};
21use crate::schema;
22
23use crate::EngineError;
24
25pub use parse::json_object;
26pub use payload::PROMPT_VERSION;
27
28pub struct GroupingOptions<'a, C: GroupingCache, A: ArtefactStore> {
29    /// The backend, always injected. `dyn` because config picks which command
30    /// to run — the one runtime-open seam in this stage (ADR 0016, 0020).
31    ///
32    /// Cancellation is a property of the backend the caller built
33    /// (`llmio::CommandBackend::with_cancel`), not of the pipeline: killing an
34    /// in-flight subprocess was never a pipeline concern.
35    pub backend: &'a dyn LlmBackend,
36    /// Where groupings are pinned. Disabling is a state of the cache
37    /// (`FsGroupingCache::disabled()`), not an `Option` here.
38    pub cache: &'a C,
39    /// Where the pre-group document is left for the model to read (ADR 0022).
40    pub artefacts: &'a A,
41    /// The executable the model runs to fetch from that document — normally
42    /// this process, so composition supplies it.
43    ///
44    /// The backend's tool allowlist MUST be built from the same string
45    /// (`llmio::CommandBackend::claude_cli`), or the prompt names a command the model
46    /// is not permitted to run.
47    pub fetch: &'a str,
48    /// Stage notifications for renderers that show progress while the
49    /// pipeline runs (the TUI's splash screen). `None` reports nothing.
50    pub progress: Option<&'a (dyn Fn(Progress) + Send + Sync)>,
51}
52
53/// Pipeline stage notifications, in the order they occur. `Grouping` carries
54/// the backend name so a renderer can say WHICH agent it is waiting on — that
55/// stage is the slow one (a subprocess LLM call on a cache miss).
56#[derive(Debug, Clone, PartialEq, Eq)]
57pub enum Progress {
58    Enumerating,
59    Classifying,
60    Grouping { backend: String, cached: bool },
61    Ordering,
62    Done,
63}
64
65/// Everything the stage needs to know about one shape class, derived from the
66/// core document + the diff view.
67pub(crate) struct ClassInfo {
68    pub id: String,
69    pub n_hunks: usize,
70    /// Exemplar hunk index into the canonical hunk list. Orders the class id
71    /// list in the prompt; the model reaches the hunk itself by fetching.
72    pub exemplar: usize,
73    /// Every member hunk lives in a generated file → noise tier.
74    pub all_generated: bool,
75    /// Some member touches a rename below the relocation threshold.
76    pub rename_gated: bool,
77    /// Sorted member digests, for the cache key.
78    pub digests: Vec<String>,
79}
80
81/// A group between audit and assembly.
82pub(crate) struct WorkGroup {
83    pub label: String,
84    pub description: String,
85    pub reason: String,
86    pub skim: bool,
87    pub class_ids: Vec<String>,
88    pub backfill: bool,
89}
90
91const RELOCATION_THRESHOLD: u8 = 95;
92
93/// Run the grouping stage over a core-only document. Returns the same document
94/// with `groups`, `reading_plan` and the grouping audit fields filled, and
95/// `"group"` appended to `generator.stages`.
96// The parameter list is the point, exactly as a bound list is: each entry is a
97// distinct authority this function may use. Bundling `langs` and `symbols`
98// behind a context struct would shorten the list without making it clearer,
99// and `CLAUDE.md` rule 2 refuses that shape.
100#[allow(clippy::too_many_arguments)]
101pub fn run<C: GroupingCache, A: ArtefactStore>(
102    doc: &schema::PlanDocument,
103    backend: &dyn LlmBackend,
104    cache: &C,
105    artefacts: &A,
106    fetch: &str,
107    lang_fingerprint: &str,
108    symbols_fingerprint: &str,
109    progress: Option<&(dyn Fn(Progress) + Send + Sync)>,
110) -> Result<schema::PlanDocument, EngineError> {
111    let infos = class_infos(doc);
112
113    let (noise, offered): (Vec<&ClassInfo>, Vec<&ClassInfo>) =
114        infos.iter().partition(|c| c.all_generated);
115
116    let mut audited = if offered.is_empty() {
117        Audited {
118            groups: Vec::new(),
119            missing: Vec::new(),
120            dupes: Vec::new(),
121            halluc: Vec::new(),
122            coverage: 1.0,
123        }
124    } else {
125        // The key names the artefact as well as the cache entry: one grouping,
126        // one document, and a cache hit finds the same file the miss wrote.
127        let key = key::cache_key(
128            &offered,
129            backend.identity(),
130            lang_fingerprint,
131            symbols_fingerprint,
132        );
133        let path = artefacts.make_readable(&key, &doc.to_json_pretty()?)?;
134        let prompt = payload::build_prompt(
135            &offered,
136            fetch,
137            &path.to_string_lossy(),
138            &doc.source.base,
139            &doc.source.head,
140        );
141        let response = fetch_response(&prompt, &key, backend, cache, progress)?;
142        let raw = parse::parse_response(&response)?;
143        audit(raw, &offered)
144    };
145
146    apply_relocation_gate(&mut audited.groups, &infos);
147
148    Ok(assemble::assemble(doc, &infos, &noise, audited))
149}
150
151pub(crate) struct Audited {
152    pub groups: Vec<WorkGroup>,
153    pub missing: Vec<String>,
154    pub dupes: Vec<String>,
155    pub halluc: Vec<String>,
156    /// Model-assigned hunks / offered hunks, pre-back-fill. The honest number.
157    pub coverage: f64,
158}
159
160/// The coverage audit — the whole point of merging class ids instead of hunk
161/// indices (ADR 0001): an omitted id is detectable and back-filled, never lost.
162fn audit(raw: parse::RawGroups, offered: &[&ClassInfo]) -> Audited {
163    let known: HashMap<&str, &ClassInfo> = offered.iter().map(|c| (c.id.as_str(), *c)).collect();
164
165    let mut claimed: HashSet<String> = HashSet::new();
166    let mut dupes = Vec::new();
167    let mut halluc = Vec::new();
168    let mut groups = Vec::new();
169
170    for g in raw.groups {
171        let mut kept = Vec::new();
172        for cid in g.classes {
173            if !known.contains_key(cid.as_str()) {
174                if !halluc.contains(&cid) {
175                    halluc.push(cid);
176                }
177            } else if claimed.contains(&cid) {
178                if !dupes.contains(&cid) {
179                    dupes.push(cid);
180                }
181            } else {
182                claimed.insert(cid.clone());
183                kept.push(cid);
184            }
185        }
186        if !kept.is_empty() {
187            groups.push(WorkGroup {
188                label: g.label,
189                description: g.description,
190                reason: g.reason,
191                skim: g.effort == "skim",
192                class_ids: kept,
193                backfill: false,
194            });
195        }
196    }
197
198    let missing: Vec<String> = offered
199        .iter()
200        .filter(|c| !claimed.contains(&c.id))
201        .map(|c| c.id.clone())
202        .collect();
203
204    let offered_hunks: usize = offered.iter().map(|c| c.n_hunks).sum();
205    let assigned_hunks: usize = offered
206        .iter()
207        .filter(|c| claimed.contains(&c.id))
208        .map(|c| c.n_hunks)
209        .sum();
210    let coverage = if offered_hunks == 0 {
211        1.0
212    } else {
213        assigned_hunks as f64 / offered_hunks as f64
214    };
215
216    if !missing.is_empty() {
217        groups.push(WorkGroup {
218            label: "Carried by no group".to_string(),
219            description: "Classes the model omitted; recovered by the coverage audit.".to_string(),
220            reason: "Not triaged — must be read.".to_string(),
221            skim: false,
222            class_ids: missing.clone(),
223            backfill: true,
224        });
225    }
226
227    Audited {
228        groups,
229        missing,
230        dupes,
231        halluc,
232        coverage,
233    }
234}
235
236/// ADR 0003: a class touching a sub-threshold rename is a modification, not a
237/// relocation, and can never stay in a skim group. Deterministic backstop that
238/// runs after the audit, whatever the model claimed.
239fn apply_relocation_gate(groups: &mut Vec<WorkGroup>, infos: &[ClassInfo]) {
240    let gated: HashSet<&str> = infos
241        .iter()
242        .filter(|c| c.rename_gated)
243        .map(|c| c.id.as_str())
244        .collect();
245    if gated.is_empty() {
246        return;
247    }
248
249    let mut extracted = Vec::new();
250    for g in groups.iter_mut() {
251        if !g.skim {
252            continue;
253        }
254        let (out, kept): (Vec<String>, Vec<String>) = g
255            .class_ids
256            .drain(..)
257            .partition(|cid| gated.contains(cid.as_str()));
258        g.class_ids = kept;
259        extracted.extend(out);
260    }
261    groups.retain(|g| !g.class_ids.is_empty());
262
263    if !extracted.is_empty() {
264        groups.push(WorkGroup {
265            label: "Modified during move".to_string(),
266            description: format!(
267                "Renamed files below the {RELOCATION_THRESHOLD}% relocation threshold: \
268                 rewritten during the move, not relocated verbatim."
269            ),
270            reason: "Rename-similarity gate: a low-similarity rename is a modification and \
271                     is never skim-eligible."
272                .to_string(),
273            skim: false,
274            class_ids: extracted,
275            backfill: false,
276        });
277    }
278}
279
280/// Derive per-class facts from the core document + view.
281fn class_infos(doc: &schema::PlanDocument) -> Vec<ClassInfo> {
282    let file_by_path: HashMap<&str, &schema::FileEntry> =
283        doc.files.iter().map(|f| (f.path.as_str(), f)).collect();
284    let generated = crate::plan::generated_files(doc);
285    let hunk_by_id: HashMap<&str, (usize, &schema::HunkEntry)> = doc
286        .hunks
287        .iter()
288        .enumerate()
289        .map(|(i, h)| (h.id.as_str(), (i, h)))
290        .collect();
291
292    doc.classes
293        .iter()
294        .map(|c| {
295            let members: Vec<(usize, &schema::HunkEntry)> = c
296                .hunk_ids
297                .iter()
298                .map(|hid| hunk_by_id[hid.as_str()])
299                .collect();
300            let mut files: Vec<String> = members.iter().map(|(_, h)| h.file.clone()).collect();
301            files.sort_unstable();
302            files.dedup();
303
304            let entries: Vec<&schema::FileEntry> =
305                files.iter().map(|p| file_by_path[p.as_str()]).collect();
306            // The one definition of the noise tier, shared with what the model
307            // is served when it asks without naming ids. Two copies of this
308            // rule would be two rules, and they would drift.
309            let all_generated = crate::plan::class_is_generated(doc, &generated, c);
310            let rename_gated = entries.iter().any(|f| {
311                f.rename_similarity
312                    .is_some_and(|s| s < RELOCATION_THRESHOLD)
313            });
314            let (exemplar_idx, _) = hunk_by_id[c.exemplar.as_str()];
315
316            let mut digests: Vec<String> = members.iter().map(|(_, h)| h.digest.clone()).collect();
317            digests.sort_unstable();
318
319            ClassInfo {
320                id: c.id.clone(),
321                n_hunks: members.len(),
322                exemplar: exemplar_idx,
323                all_generated,
324                rename_gated,
325                digests,
326            }
327        })
328        .collect()
329}
330
331/// Cache-or-call: the cached value is the raw model response, so the audit and
332/// assembly stay pure functions replayed on every load.
333fn fetch_response<C: GroupingCache>(
334    prompt: &str,
335    key: &str,
336    backend: &dyn LlmBackend,
337    cache: &C,
338    progress: Option<&(dyn Fn(Progress) + Send + Sync)>,
339) -> Result<String, EngineError> {
340    let report = |cached: bool| {
341        if let Some(f) = progress {
342            f(Progress::Grouping {
343                backend: backend.name().to_string(),
344                cached,
345            });
346        }
347    };
348    if let Some(hit) = cache.get(key)? {
349        report(true);
350        return Ok(hit);
351    }
352    report(false);
353    let response = backend.complete(prompt)?;
354    cache.put(key, &response)?;
355    Ok(response)
356}