Move scaffold/scaffold/discipline_gate.py -> scaffold/discipline_gate.py

scaffold/discipline_gate.py

@@ -0,0 +1,175 @@
+"""Rivet Discipline Gate — pre-flight check on every suggestion.
+
+Before any code suggestion reaches the user, it passes through this gate.
+The gate checks for dangerous patterns, flags security implications, and
+assigns a confidence level. If a suggestion fails the gate, it's blocked
+or annotated with warnings.
+"""
+
+import re
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class Confidence(str, Enum):
+    HIGH = "HIGH"      # Traced the full code path
+    MEDIUM = "MEDIUM"  # Read the code but not run it
+    LOW = "LOW"        # Reasoning from architecture, not source
+
+
+class Severity(str, Enum):
+    BLOCK = "BLOCK"    # Do not suggest this
+    WARN = "WARN"      # Suggest with prominent warning
+    INFO = "INFO"      # Note for awareness
+
+
+@dataclass
+class GateResult:
+    passed: bool = True
+    flags: list = field(default_factory=list)
+    confidence: Confidence = Confidence.MEDIUM
+    requires_review: list = field(default_factory=list)
+
+    def add_flag(self, severity: Severity, message: str):
+        self.flags.append({"severity": severity.value, "message": message})
+        if severity == Severity.BLOCK:
+            self.passed = False
+
+    def format_warnings(self) -> str:
+        if not self.flags:
+            return ""
+        lines = ["\n⚠️ **Discipline Gate Flags:**"]
+        for f in self.flags:
+            icon = "🛑" if f["severity"] == "BLOCK" else "⚠️" if f["severity"] == "WARN" else "ℹ️"
+            lines.append(f"  {icon} [{f['severity']}] {f['message']}")
+        if self.requires_review:
+            lines.append(f"\n  📋 Requires review from: {', '.join(self.requires_review)}")
+        lines.append(f"\n  📊 Confidence: {self.confidence.value}")
+        return "\n".join(lines)
+
+
+# Pattern matchers for dangerous operations
+DESTRUCTIVE_MIGRATION_PATTERNS = [
+    r'\bDROP\s+(TABLE|COLUMN|INDEX|CONSTRAINT)',
+    r'\bALTER\s+TABLE\s+\w+\s+DROP',
+    r'\bALTER\s+TABLE\s+\w+\s+ALTER\s+COLUMN\s+\w+\s+TYPE',
+    r'\bTRUNCATE\s+TABLE',
+    r'\bDELETE\s+FROM\s+\w+\s*;',  # Unqualified DELETE
+]
+
+AUTH_PATTERNS = [
+    r'\bJWT_SECRET\b',
+    r'\bverifyToken\b',
+    r'\brequireAdmin\b',
+    r'\brequireModerator\b',
+    r'\bauth\s*middleware\b',
+    r'\bsession\s*cookie\b',
+    r'\brejectIfIneligible\b',
+    r'\bBearer\b',
+]
+
+WEBHOOK_BYPASS_PATTERNS = [
+    r"process\.env\.\w+\s*\|\|\s*['\"]",  # env fallback to empty string
+    r"if\s*\(\s*!secret\s*\)\s*return",   # skip on missing secret
+    r"if\s*\(\s*!.*SECRET.*\)\s*return",
+]
+
+TRANSACTION_GAPS = [
+    r'await\s+\w+\.(debit|credit|transfer|purchase|delete)\(',  # Multi-step without tx
+]
+
+
+def check_migration_safety(content: str) -> list:
+    """Check for destructive migration patterns."""
+    flags = []
+    for pattern in DESTRUCTIVE_MIGRATION_PATTERNS:
+        matches = re.findall(pattern, content, re.IGNORECASE)
+        if matches:
+            flags.append((Severity.BLOCK,
+                         f"Destructive migration detected: {pattern}. "
+                         f"Staging and prod share the database — additive only."))
+    return flags
+
+
+def check_auth_impact(content: str) -> list:
+    """Flag anything touching authentication."""
+    flags = []
+    for pattern in AUTH_PATTERNS:
+        if re.search(pattern, content, re.IGNORECASE):
+            flags.append((Severity.WARN,
+                         "This touches authentication. Review with security before merging."))
+            break
+    return flags
+
+
+def check_webhook_safety(content: str) -> list:
+    """Check for webhook signature bypass patterns (Audit C2)."""
+    flags = []
+    for pattern in WEBHOOK_BYPASS_PATTERNS:
+        if re.search(pattern, content):
+            flags.append((Severity.WARN,
+                         "Webhook signature bypass pattern detected (Audit C2). "
+                         "Never skip verification when env var is missing — reject instead."))
+            break
+    return flags
+
+
+def check_transaction_safety(content: str) -> list:
+    """Flag multi-step DB operations that may need transaction wrapping (Audit H3)."""
+    flags = []
+    matches = re.findall(r'await\s+\w+\.\w+\(', content)
+    if len(matches) >= 2:
+        for pattern in TRANSACTION_GAPS:
+            if re.search(pattern, content):
+                flags.append((Severity.WARN,
+                             "Multi-step DB operation detected. Ensure these are wrapped "
+                             "in a single transaction (Audit H3: gem debit without tx protection)."))
+                break
+    return flags
+
+
+def run_gate(suggestion: str, context: str = "") -> GateResult:
+    """Run the full discipline gate on a suggestion.
+
+    Args:
+        suggestion: The code/text being suggested to the user
+        context: Optional surrounding context (the question, file being edited)
+
+    Returns:
+        GateResult with pass/fail, flags, and confidence
+    """
+    result = GateResult()
+    full_content = f"{context}\n{suggestion}"
+
+    # Run all checks
+    for severity, message in check_migration_safety(full_content):
+        result.add_flag(severity, message)
+
+    for severity, message in check_auth_impact(full_content):
+        result.add_flag(severity, message)
+        result.requires_review.append("security")
+
+    for severity, message in check_webhook_safety(full_content):
+        result.add_flag(severity, message)
+        result.requires_review.append("security")
+
+    for severity, message in check_transaction_safety(full_content):
+        result.add_flag(severity, message)
+
+    return result
+
+
+if __name__ == "__main__":
+    # Quick test
+    test_migration = "ALTER TABLE students DROP COLUMN legacy_score;"
+    result = run_gate(test_migration)
+    print(f"Migration test: passed={result.passed}")
+    print(result.format_warnings())
+
+    test_webhook = """
+    const secret = process.env.STRIPE_SECRET || ''
+    if (!secret) return next()
+    """
+    result = run_gate(test_webhook)
+    print(f"\nWebhook test: passed={result.passed}")
+    print(result.format_warnings())