forge_guard/ai/

consensus.rs

//! Consensus engine — aggregates findings from multiple AI auditors with
//! cross-validation confidence scoring and deduplication.

use crate::ai::auditors::AuditorAgent;
use crate::ai::{AuditContext, AuditorFinding};

use super::ConsensusFinding;

/// Minimum confidence threshold for a finding to be included in the final report.
const DEFAULT_MIN_CONFIDENCE: f64 = 0.5;

/// Consensus configuration.
#[derive(Debug, Clone)]
pub struct ConsensusConfig {
    /// Minimum confidence (0.0–1.0) required for a finding to be included.
    pub min_confidence: f64,
    /// Boost confidence when multiple auditors flag the same issue.
    pub consensus_boost: f64,
    /// Whether to merge findings that match across domains.
    pub cross_domain_merge: bool,
}

impl Default for ConsensusConfig {
    fn default() -> Self {
        Self {
            min_confidence: DEFAULT_MIN_CONFIDENCE,
            consensus_boost: 0.2,
            cross_domain_merge: true,
        }
    }
}

/// A report produced by the consensus engine.
#[derive(Debug, Clone)]
pub struct ConsensusReport {
    /// All findings that passed the consensus threshold.
    pub findings: Vec<ConsensusFinding>,
    /// Number of auditors that participated.
    pub auditor_count: usize,
    /// How many findings were deduplicated / merged.
    pub deduplicated_count: usize,
    /// How many findings were discarded due to low confidence.
    pub filtered_count: usize,
}

/// The consensus engine runs multiple auditors over the same source code,
/// deduplicates overlapping findings, and boosts confidence when multiple
/// auditors independently flag the same issue.
pub struct ConsensusEngine {
    auditors: Vec<Box<dyn AuditorAgent>>,
    config: ConsensusConfig,
}

impl ConsensusEngine {
    /// Create a new consensus engine with the given configuration.
    pub fn new(config: ConsensusConfig) -> Self {
        Self {
            auditors: Vec::new(),
            config,
        }
    }

    /// Register an auditor agent.
    pub fn register(&mut self, auditor: Box<dyn AuditorAgent>) {
        self.auditors.push(auditor);
    }

    /// Number of registered auditors.
    pub fn auditor_count(&self) -> usize {
        self.auditors.len()
    }

    /// Run all auditors and aggregate findings through the consensus pipeline.
    pub fn analyze(&self, context: &AuditContext) -> ConsensusReport {
        let mut all_findings: Vec<ConsensusFinding> = Vec::new();

        for auditor in &self.auditors {
            match auditor.analyze(context) {
                Ok(findings) => {
                    for finding in findings {
                        all_findings.push(ConsensusFinding {
                            auditor: auditor.name().to_owned(),
                            domain: auditor.domain().to_owned(),
                            finding,
                            cross_validated: false,
                        });
                    }
                }
                Err(e) => {
                    eprintln!("  ⚠️  [ai:{}] Error: {e}", auditor.name());
                }
            }
        }

        // Step 1: Cross-validate — boost confidence for duplicates across auditors
        let (mut merged, dedup_count) = self.deduplicate_and_boost(all_findings);

        // Step 2: Filter by minimum confidence
        let before_filter = merged.len();
        merged.retain(|cf| cf.finding.confidence >= self.config.min_confidence);
        let filtered = before_filter - merged.len();

        ConsensusReport {
            findings: merged,
            auditor_count: self.auditors.len(),
            deduplicated_count: dedup_count,
            filtered_count: filtered,
        }
    }

    /// Deduplicate findings across auditors and boost confidence for consensus.
    ///
    /// Two findings are considered "the same issue" when their normalized titles
    /// share significant keyword overlap. When multiple auditors flag the same
    /// issue, confidence is boosted by `consensus_boost` per additional auditor.
    fn deduplicate_and_boost(
        &self,
        findings: Vec<ConsensusFinding>,
    ) -> (Vec<ConsensusFinding>, usize) {
        if findings.is_empty() {
            return (Vec::new(), 0);
        }

        let mut dedup_count = 0;
        let mut groups: Vec<Vec<ConsensusFinding>> = Vec::new();

        for finding in findings {
            // Try to match into an existing group
            let matched = if self.config.cross_domain_merge {
                groups.iter_mut().find(|group| {
                    group
                        .first()
                        .map(|first| same_issue(&first.finding, &finding.finding))
                        .unwrap_or(false)
                })
            } else {
                // Only match within the same domain
                groups.iter_mut().find(|group| {
                    group
                        .first()
                        .map(|first| {
                            first.domain == finding.domain
                                && same_issue(&first.finding, &finding.finding)
                        })
                        .unwrap_or(false)
                })
            };

            match matched {
                Some(group) => {
                    group.push(finding);
                    dedup_count += 1;
                }
                None => {
                    groups.push(vec![finding]);
                }
            }
        }

        // Apply consensus boost and keep the highest-confidence finding per group
        let mut result = Vec::with_capacity(groups.len());
        for group in groups {
            let count = group.len();
            let boost = if count > 1 {
                (count as f64 - 1.0) * self.config.consensus_boost
            } else {
                0.0
            };

            // Find the finding with highest original confidence
            let mut best = group.into_iter().max_by(|a, b| {
                a.finding
                    .confidence
                    .partial_cmp(&b.finding.confidence)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });

            if let Some(ref mut best) = best {
                best.finding.confidence = (best.finding.confidence + boost).min(1.0);
                best.cross_validated = count > 1;
            }

            if let Some(best) = best {
                result.push(best);
            }
        }

        (result, dedup_count)
    }
}

/// Determine whether two findings describe the same underlying issue.
///
/// Uses title keyword overlap as a heuristic:
/// - Tokenize both titles into lowercase keywords
/// - If intersection / min(|a|, |b|) >= 0.5, they're the same issue
fn same_issue(a: &AuditorFinding, b: &AuditorFinding) -> bool {
    let tokens_a = tokenize(&a.title);
    let tokens_b = tokenize(&b.title);

    if tokens_a.is_empty() || tokens_b.is_empty() {
        return false;
    }

    let intersection: Vec<&String> = tokens_a.iter().filter(|t| tokens_b.contains(t)).collect();
    let min_len = tokens_a.len().min(tokens_b.len());
    (intersection.len() as f64 / min_len as f64) >= 0.5
}

/// Normalize a string into lowercase keyword tokens.
fn tokenize(s: &str) -> Vec<String> {
    s.to_lowercase()
        .split(|c: char| !c.is_alphanumeric())
        .filter(|t| !t.is_empty() && t.len() > 2)
        .map(String::from)
        .collect()
}

// ── Tests ────────────────────────────────────────────────────
#[cfg(test)]
mod tests {
    use super::*;
    use crate::core::Severity;

    fn make_finding(
        title: &str,
        confidence: f64,
        severity: Severity,
        category: &str,
    ) -> AuditorFinding {
        AuditorFinding {
            title: title.to_owned(),
            description: "description".into(),
            confidence,
            severity,
            suggestion: "recommendation".into(),
            line_numbers: vec![1],
            category: category.to_owned(),
        }
    }

    fn make_consensus(
        auditor: &str,
        domain: &str,
        title: &str,
        confidence: f64,
    ) -> ConsensusFinding {
        ConsensusFinding {
            auditor: auditor.to_owned(),
            domain: domain.to_owned(),
            finding: make_finding(title, confidence, Severity::High, "Test"),
            cross_validated: false,
        }
    }

    #[test]
    fn test_empty_engine() {
        let engine = ConsensusEngine::new(ConsensusConfig::default());
        let report = engine.analyze(&AuditContext {
            source_code: "".into(),
            file_name: "test.sol".into(),
            compiler_version: "0.8.20".into(),
            additional: Default::default(),
        });
        assert_eq!(report.auditor_count, 0);
        assert!(report.findings.is_empty());
    }

    #[test]
    fn test_tokenize() {
        let tokens = tokenize("Reentrancy vulnerability in withdraw()");
        assert!(tokens.contains(&"reentrancy".to_string()));
        assert!(tokens.contains(&"vulnerability".to_string()));
        assert!(tokens.contains(&"withdraw".to_string()));
        assert!(!tokens.contains(&"in".to_string())); // too short
    }

    #[test]
    fn test_same_issue_identical() {
        let a = make_finding("Reentrancy in withdraw", 0.9, Severity::High, "Reentrancy");
        let b = make_finding("Reentrancy in withdraw", 0.8, Severity::High, "Reentrancy");
        assert!(same_issue(&a, &b));
    }

    #[test]
    fn test_same_issue_similar() {
        let a = make_finding(
            "Reentrancy vulnerability in withdraw function",
            0.9,
            Severity::High,
            "Reentrancy",
        );
        let b = make_finding(
            "Reentrancy bug in withdraw",
            0.8,
            Severity::High,
            "Reentrancy",
        );
        assert!(same_issue(&a, &b));
    }

    #[test]
    fn test_different_issues() {
        let a = make_finding("Reentrancy in withdraw", 0.9, Severity::High, "Reentrancy");
        let b = make_finding(
            "Missing access control on mint",
            0.8,
            Severity::High,
            "AccessControl",
        );
        assert!(!same_issue(&a, &b));
    }

    #[test]
    fn test_deduplicate_and_boost() {
        let engine = ConsensusEngine::new(ConsensusConfig::default());
        let findings = vec![
            make_consensus("auditor-a", "Security", "Reentrancy in withdraw", 0.8),
            make_consensus("auditor-b", "Security", "Reentrancy in withdraw", 0.7),
            make_consensus("auditor-a", "Security", "Access control on mint", 0.9),
        ];

        let (result, dedup) = engine.deduplicate_and_boost(findings);
        assert_eq!(
            result.len(),
            2,
            "should merge two reentrancy findings into one"
        );
        assert_eq!(dedup, 1, "one finding should be deduplicated");

        // The surviving reentrancy finding should have boosted confidence
        let reentrancy = result
            .iter()
            .find(|cf| cf.finding.title.contains("Reentrancy"))
            .unwrap();
        assert!(
            reentrancy.finding.confidence > 0.9,
            "confidence should be boosted: {}",
            reentrancy.finding.confidence
        );
        assert!(reentrancy.cross_validated);
    }

    #[test]
    fn test_no_deduplication_when_different_titles() {
        let engine = ConsensusEngine::new(ConsensusConfig::default());
        let findings = vec![
            make_consensus("auditor-a", "Security", "Reentrancy in withdraw", 0.8),
            make_consensus("auditor-b", "Security", "Access control on mint", 0.7),
            make_consensus("auditor-c", "Gas", "Loop gas waste", 0.6),
        ];

        let (result, dedup) = engine.deduplicate_and_boost(findings);
        assert_eq!(result.len(), 3);
        assert_eq!(dedup, 0);
    }

    #[test]
    fn test_filter_low_confidence() {
        let engine = ConsensusEngine::new(ConsensusConfig {
            min_confidence: 0.7,
            ..Default::default()
        });

        // We need at least one auditor registered for the report
        // Just test the filtering manually
        let findings = vec![
            make_consensus("auditor-a", "Security", "Critical reentrancy", 0.95),
            make_consensus("auditor-b", "Security", "Low confidence issue", 0.3),
        ];

        // Simulate what analyze() does
        let (mut merged, _) = engine.deduplicate_and_boost(findings);
        let before = merged.len();
        merged.retain(|cf| cf.finding.confidence >= 0.7);
        let filtered = before - merged.len();

        assert_eq!(merged.len(), 1);
        assert_eq!(filtered, 1);
        assert!(merged[0].finding.title.contains("reentrancy"));
    }

    #[test]
    fn test_consensus_report_structure() {
        // Test with no auditors registered
        let engine = ConsensusEngine::new(ConsensusConfig::default());
        let ctx = AuditContext {
            source_code: "contract C {}".into(),
            file_name: "c.sol".into(),
            compiler_version: "0.8.20".into(),
            additional: Default::default(),
        };
        let report = engine.analyze(&ctx);
        assert_eq!(report.auditor_count, 0);
        assert_eq!(report.deduplicated_count, 0);
        assert_eq!(report.filtered_count, 0);
    }
}