compliance-scanner-agent/compliance-agent/src/controls/semantic.rs

//! Semantic control mapping: retrieve the top-K controls nearest a code region,
//! then confirm each with the grounded judge.
//!
//! The `region → controls` direction (vs. the CWE-LUT's `finding → control`) is
//! what scales to the full master-control corpus: the LLM only ever judges a
//! handful of retrieved candidates, and every surviving verdict is still anchored
//! to real code by the grounding gate.

use compliance_core::control_check::{ground, CandidateRegion};
use compliance_core::models::Finding;

use super::index::ControlIndex;
use super::judge::ControlJudge;

/// Retrieve → judge → ground, generic over the judge so tests use a stub.
pub struct SemanticControlChecker<J> {
    judge: J,
}

impl<J: ControlJudge> SemanticControlChecker<J> {
    pub fn new(judge: J) -> Self {
        Self { judge }
    }

    /// Map a code region to the controls it violates. `query_embedding` is the
    /// caller-supplied retrieval embedding — typically the finding's intent
    /// (title/description) plus the region, so retrieval keys on what the finding
    /// is *about*, not just the ambient code. The top-`k` nearest controls in
    /// `index` are then judged against the raw `region` and grounded.
    pub async fn check(
        &self,
        index: &ControlIndex,
        region: &CandidateRegion,
        query_embedding: &[f64],
        k: usize,
        repo_id: &str,
    ) -> Vec<Finding> {
        let candidates = index.nearest(query_embedding, k);
        let mut findings = Vec::new();
        for spec in &candidates {
            let verdict = self.judge.judge(spec, region).await;
            if let Some(finding) = ground(spec, region, &verdict, repo_id) {
                findings.push(finding);
            }
        }
        findings
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use compliance_core::control_check::{ControlCheckSpec, LlmVerdict};
    use compliance_core::models::finding::Severity;

    struct StubJudge {
        verdict: LlmVerdict,
    }
    impl ControlJudge for StubJudge {
        async fn judge(&self, _s: &ControlCheckSpec, _r: &CandidateRegion) -> LlmVerdict {
            self.verdict.clone()
        }
    }

    fn spec(id: &str) -> ControlCheckSpec {
        ControlCheckSpec {
            control_id: id.into(),
            title: id.into(),
            requirement: id.into(),
            default_cwe: None,
            severity: Severity::Medium,
        }
    }

    #[tokio::test]
    async fn retrieves_then_grounds_the_nearest_control() {
        let index = ControlIndex::from_embeddings(vec![
            (spec("mc-near"), vec![1.0, 0.0]),
            (spec("mc-far"), vec![0.0, 1.0]),
        ]);
        let checker = SemanticControlChecker::new(StubJudge {
            verdict: LlmVerdict {
                violates: true,
                snippet: "PASSWORD = \"admin\"".into(),
                cwe: None,
                confidence: 0.9,
            },
        });
        let region = CandidateRegion {
            file: "src/auth.py".into(),
            start_line: 1,
            content: "PASSWORD = \"admin\"\n".into(),
        };
        // Query embedding nearest to mc-near; k=1 → only mc-near is judged.
        let findings = checker
            .check(&index, &region, &[0.95, 0.05], 1, "repo")
            .await;
        assert_eq!(findings.len(), 1);
        assert_eq!(findings[0].control_refs, vec!["mc-near".to_string()]);
    }

    #[tokio::test]
    async fn ungrounded_verdict_is_dropped() {
        let index = ControlIndex::from_embeddings(vec![(spec("mc-near"), vec![1.0, 0.0])]);
        let checker = SemanticControlChecker::new(StubJudge {
            verdict: LlmVerdict {
                violates: true,
                snippet: "not in the region".into(),
                cwe: None,
                confidence: 0.9,
            },
        });
        let region = CandidateRegion {
            file: "f".into(),
            start_line: 1,
            content: "real code\n".into(),
        };
        let findings = checker.check(&index, &region, &[1.0, 0.0], 1, "repo").await;
        assert!(findings.is_empty());
    }
}