crucible-mcp 0.3.0__py3-none-any.whl → 0.5.0__py3-none-any.whl

crucible/enforcement/__init__.py

@@ -0,0 +1,40 @@
+"""Enforcement module for pattern assertions and applicability checking."""
+
+from crucible.enforcement.assertions import (
+    get_all_assertion_files,
+    load_assertion_file,
+    load_assertions,
+    resolve_assertion_file,
+)
+from crucible.enforcement.models import (
+    Assertion,
+    AssertionFile,
+    AssertionType,
+    PatternMatch,
+    Priority,
+    Suppression,
+)
+from crucible.enforcement.patterns import (
+    find_pattern_matches,
+    parse_suppressions,
+    run_pattern_assertions,
+)
+
+__all__ = [
+    # Models
+    "Assertion",
+    "AssertionFile",
+    "AssertionType",
+    "PatternMatch",
+    "Priority",
+    "Suppression",
+    # Assertions
+    "get_all_assertion_files",
+    "load_assertion_file",
+    "load_assertions",
+    "resolve_assertion_file",
+    # Patterns
+    "find_pattern_matches",
+    "parse_suppressions",
+    "run_pattern_assertions",
+]

crucible/enforcement/assertions.py

@@ -0,0 +1,276 @@
+"""Load and validate assertion files.
+
+Assertions follow the same cascade as skills/knowledge:
+1. Project: .crucible/assertions/
+2. User: ~/.claude/crucible/assertions/
+3. Bundled: package assertions/ (none bundled by default)
+"""
+
+from functools import lru_cache
+from pathlib import Path
+
+import yaml
+
+from crucible.enforcement.models import (
+    Applicability,
+    Assertion,
+    AssertionFile,
+    AssertionType,
+    Priority,
+)
+from crucible.errors import Result, err, ok
+
+# Assertion directories (cascade priority)
+ASSERTIONS_BUNDLED = Path(__file__).parent / "bundled"
+ASSERTIONS_USER = Path.home() / ".claude" / "crucible" / "assertions"
+ASSERTIONS_PROJECT = Path(".crucible") / "assertions"
+
+
+def resolve_assertion_file(filename: str) -> tuple[Path | None, str]:
+    """Find assertion file with cascade priority.
+
+    Returns (path, source) where source is 'project', 'user', or 'bundled'.
+    """
+    # Ensure .yaml extension
+    if not filename.endswith((".yaml", ".yml")):
+        filename = f"{filename}.yaml"
+
+    # 1. Project-level (highest priority)
+    project_path = ASSERTIONS_PROJECT / filename
+    if project_path.exists():
+        return project_path, "project"
+
+    # 2. User-level
+    user_path = ASSERTIONS_USER / filename
+    if user_path.exists():
+        return user_path, "user"
+
+    # 3. Bundled (lowest priority)
+    bundled_path = ASSERTIONS_BUNDLED / filename
+    if bundled_path.exists():
+        return bundled_path, "bundled"
+
+    return None, ""
+
+
+def get_all_assertion_files() -> set[str]:
+    """Get all available assertion file names from all sources."""
+    files: set[str] = set()
+
+    for source_dir in [ASSERTIONS_BUNDLED, ASSERTIONS_USER, ASSERTIONS_PROJECT]:
+        if source_dir.exists():
+            for file_path in source_dir.iterdir():
+                if file_path.is_file() and file_path.suffix in (".yaml", ".yml"):
+                    files.add(file_path.name)
+
+    return files
+
+
+def _parse_priority(value: str | None) -> Priority:
+    """Parse priority string to enum."""
+    if value is None:
+        return Priority.MEDIUM
+    try:
+        return Priority(value.lower())
+    except ValueError:
+        return Priority.MEDIUM
+
+
+def _parse_applicability(data: dict | None) -> Applicability | None:
+    """Parse applicability configuration."""
+    if data is None:
+        return None
+
+    exclude_raw = data.get("exclude", [])
+    if isinstance(exclude_raw, str):
+        exclude_raw = [exclude_raw]
+
+    return Applicability(
+        glob=data.get("glob"),
+        exclude=tuple(exclude_raw),
+    )
+
+
+def _parse_assertion(data: dict) -> Result[Assertion, str]:
+    """Parse a single assertion from YAML data."""
+    # Required fields
+    assertion_id = data.get("id")
+    if not assertion_id:
+        return err("Assertion missing required 'id' field")
+
+    type_str = data.get("type", "pattern")
+    try:
+        assertion_type = AssertionType(type_str)
+    except ValueError:
+        return err(f"Invalid assertion type: {type_str}")
+
+    message = data.get("message")
+    if not message:
+        return err(f"Assertion '{assertion_id}' missing required 'message' field")
+
+    # Optional fields with defaults
+    severity = data.get("severity", "warning")
+    if severity not in ("error", "warning", "info"):
+        severity = "warning"
+
+    priority = _parse_priority(data.get("priority"))
+
+    # Pattern-specific fields
+    pattern = data.get("pattern")
+    if assertion_type == AssertionType.PATTERN and not pattern:
+        return err(f"Pattern assertion '{assertion_id}' missing required 'pattern' field")
+
+    # Languages
+    languages_raw = data.get("languages", [])
+    if isinstance(languages_raw, str):
+        languages_raw = [languages_raw]
+    languages = tuple(lang.lower() for lang in languages_raw)
+
+    # Applicability
+    applicability = _parse_applicability(data.get("applicability"))
+
+    # LLM-specific fields (v0.5+)
+    compliance = data.get("compliance")
+    model = data.get("model")
+
+    return ok(
+        Assertion(
+            id=assertion_id,
+            type=assertion_type,
+            message=message,
+            severity=severity,  # type: ignore[arg-type]
+            priority=priority,
+            pattern=pattern,
+            languages=languages,
+            applicability=applicability,
+            compliance=compliance,
+            model=model,
+        )
+    )
+
+
+def _validate_assertion_file(data: dict, path: str) -> Result[AssertionFile, str]:
+    """Validate and parse an assertion file."""
+    version = data.get("version", "0.4")
+    name = data.get("name", "")
+    description = data.get("description", "")
+
+    assertions_data = data.get("assertions", [])
+    if not isinstance(assertions_data, list):
+        return err(f"{path}: 'assertions' must be a list")
+
+    assertions: list[Assertion] = []
+    for i, assertion_data in enumerate(assertions_data):
+        if not isinstance(assertion_data, dict):
+            return err(f"{path}: assertion {i} must be an object")
+
+        result = _parse_assertion(assertion_data)
+        if result.is_err:
+            return err(f"{path}: {result.error}")
+        assertions.append(result.value)
+
+    # Check for duplicate IDs
+    seen_ids: set[str] = set()
+    for assertion in assertions:
+        if assertion.id in seen_ids:
+            return err(f"{path}: duplicate assertion ID '{assertion.id}'")
+        seen_ids.add(assertion.id)
+
+    return ok(
+        AssertionFile(
+            version=str(version),
+            name=name,
+            description=description,
+            assertions=tuple(assertions),
+            source="",  # Set by caller
+            path=path,
+        )
+    )
+
+
+@lru_cache(maxsize=64)
+def _load_assertion_file_cached(path_str: str) -> AssertionFile | str:
+    """Internal cached assertion file loader.
+
+    Returns AssertionFile on success, error string on failure.
+    """
+    path = Path(path_str)
+    try:
+        content = path.read_text()
+        data = yaml.safe_load(content)
+    except OSError as e:
+        return f"Failed to read '{path}': {e}"
+    except yaml.YAMLError as e:
+        return f"Invalid YAML in '{path}': {e}"
+
+    if not isinstance(data, dict):
+        return f"'{path}' must contain a YAML object"
+
+    result = _validate_assertion_file(data, str(path))
+    if result.is_err:
+        return result.error
+
+    return result.value
+
+
+def load_assertion_file(filename: str) -> Result[AssertionFile, str]:
+    """Load a single assertion file by name with cascade resolution.
+
+    Args:
+        filename: Assertion file name (e.g., "security.yaml")
+
+    Returns:
+        Result containing AssertionFile or error message
+    """
+    path, source = resolve_assertion_file(filename)
+    if path is None:
+        return err(f"Assertion file '{filename}' not found")
+
+    cached = _load_assertion_file_cached(str(path))
+    if isinstance(cached, str):
+        return err(cached)
+
+    # Return a new AssertionFile with the correct source
+    return ok(
+        AssertionFile(
+            version=cached.version,
+            name=cached.name,
+            description=cached.description,
+            assertions=cached.assertions,
+            source=source,
+            path=str(path),
+        )
+    )
+
+
+def load_assertions(filenames: set[str] | None = None) -> tuple[list[Assertion], list[str]]:
+    """Load all assertions from specified or all available files.
+
+    Args:
+        filenames: Specific files to load (if None, loads all)
+
+    Returns:
+        Tuple of (list of assertions, list of error messages)
+    """
+    if filenames is None:
+        filenames = get_all_assertion_files()
+
+    assertions: list[Assertion] = []
+    errors: list[str] = []
+
+    for filename in sorted(filenames):
+        result = load_assertion_file(filename)
+        if result.is_err:
+            errors.append(result.error)
+        else:
+            assertions.extend(result.value.assertions)
+
+    # Sort by priority (critical first)
+    assertions.sort(key=lambda a: a.priority.rank)
+
+    return assertions, errors
+
+
+def clear_assertion_cache() -> None:
+    """Clear the assertion loading cache. Useful for testing or after updates."""
+    _load_assertion_file_cached.cache_clear()

crucible/enforcement/budget.py

@@ -0,0 +1,179 @@
+"""Token budget estimation and tracking for LLM compliance assertions."""
+
+from crucible.enforcement.models import (
+    Assertion,
+    AssertionType,
+    BudgetState,
+    ComplianceConfig,
+)
+
+# Average tokens per character (rough estimate for code)
+TOKENS_PER_CHAR = 0.25
+
+# Base overhead for each LLM call (system prompt, response format, etc.)
+BASE_OVERHEAD_TOKENS = 200
+
+# Minimum tokens for compliance prompt
+MIN_COMPLIANCE_TOKENS = 50
+
+
+def estimate_assertion_tokens(assertion: Assertion, content_length: int) -> int:
+    """Estimate tokens needed to run an LLM assertion.
+
+    Args:
+        assertion: The assertion to estimate
+        content_length: Length of content to analyze in characters
+
+    Returns:
+        Estimated token count for input
+    """
+    if assertion.type != AssertionType.LLM:
+        return 0
+
+    # Content tokens
+    content_tokens = int(content_length * TOKENS_PER_CHAR)
+
+    # Compliance prompt tokens
+    compliance_tokens = 0
+    if assertion.compliance:
+        compliance_tokens = max(
+            MIN_COMPLIANCE_TOKENS,
+            int(len(assertion.compliance) * TOKENS_PER_CHAR),
+        )
+
+    return BASE_OVERHEAD_TOKENS + content_tokens + compliance_tokens
+
+
+def estimate_total_budget(
+    assertions: list[Assertion],
+    content_length: int,
+) -> int:
+    """Estimate total tokens needed to run all LLM assertions.
+
+    Args:
+        assertions: List of assertions (filters to LLM only)
+        content_length: Length of content to analyze
+
+    Returns:
+        Estimated total token count
+    """
+    total = 0
+    for assertion in assertions:
+        if assertion.type == AssertionType.LLM:
+            total += estimate_assertion_tokens(assertion, content_length)
+    return total
+
+
+def sort_by_priority(assertions: list[Assertion]) -> list[Assertion]:
+    """Sort assertions by priority (critical first).
+
+    Args:
+        assertions: Assertions to sort
+
+    Returns:
+        Sorted list (critical > high > medium > low)
+    """
+    return sorted(assertions, key=lambda a: a.priority.rank)
+
+
+def select_within_budget(
+    assertions: list[Assertion],
+    content_length: int,
+    budget: int,
+) -> tuple[list[Assertion], list[Assertion]]:
+    """Select assertions that fit within token budget.
+
+    Args:
+        assertions: Assertions to select from (should be pre-sorted by priority)
+        content_length: Length of content to analyze
+        budget: Token budget (0 = unlimited)
+
+    Returns:
+        Tuple of (selected_assertions, skipped_assertions)
+    """
+    if budget == 0:
+        # Unlimited budget
+        return list(assertions), []
+
+    selected: list[Assertion] = []
+    skipped: list[Assertion] = []
+    tokens_used = 0
+
+    for assertion in assertions:
+        if assertion.type != AssertionType.LLM:
+            continue
+
+        estimated = estimate_assertion_tokens(assertion, content_length)
+
+        if tokens_used + estimated <= budget:
+            selected.append(assertion)
+            tokens_used += estimated
+        else:
+            skipped.append(assertion)
+
+    return selected, skipped
+
+
+def filter_llm_assertions(assertions: list[Assertion]) -> list[Assertion]:
+    """Filter to only LLM-type assertions.
+
+    Args:
+        assertions: All assertions
+
+    Returns:
+        Only assertions with type=llm
+    """
+    return [a for a in assertions if a.type == AssertionType.LLM]
+
+
+def create_budget_state(config: ComplianceConfig) -> BudgetState:
+    """Create initial budget state from config.
+
+    Args:
+        config: Compliance configuration
+
+    Returns:
+        Fresh BudgetState
+    """
+    return BudgetState(total_budget=config.token_budget)
+
+
+def prepare_llm_assertions(
+    assertions: list[Assertion],
+    content_length: int,
+    config: ComplianceConfig,
+) -> tuple[list[Assertion], BudgetState]:
+    """Prepare LLM assertions for execution.
+
+    Filters to LLM assertions, sorts by priority, and selects within budget.
+
+    Args:
+        assertions: All loaded assertions
+        content_length: Length of content to analyze
+        config: Compliance configuration
+
+    Returns:
+        Tuple of (assertions_to_run, budget_state)
+    """
+    # Filter to LLM assertions only
+    llm_assertions = filter_llm_assertions(assertions)
+
+    if not llm_assertions:
+        return [], create_budget_state(config)
+
+    # Sort by priority
+    sorted_assertions = sort_by_priority(llm_assertions)
+
+    # Select within budget
+    selected, skipped = select_within_budget(
+        sorted_assertions,
+        content_length,
+        config.token_budget,
+    )
+
+    # Create budget state
+    state = create_budget_state(config)
+    for assertion in skipped:
+        state.skip(assertion.id)
+
+    return selected, state