traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/governance/registry.py

@@ -0,0 +1,110 @@
+"""Per-session residency for governance: the single owner of live SessionState.
+
+A session has two distinct residency scopes that must never be conflated, because
+the two governance channels have opposite durability and threading constraints:
+
+  * **Durable / observation residency** — owned by the ``SessionMonitor`` single
+    writer. States are rehydrated from the durable store, carry a live DB
+    connection, and their ``persist`` actually writes. Phase-1 mutations here
+    (budget, taint) must survive a restart.
+  * **Ephemeral / gate residency** — owned by the ``Shield`` enforcement gate,
+    which runs at the framework's input edge on arbitrary threads. Gate state is
+    created with ``_db=None`` and never touches SQLite, so it is safe to build
+    cross-thread; it only carries in-process gate context.
+
+Keeping these separate is the fix for the F2 durability gap: when the gate creates
+a session's state first, the writer must not be handed that non-durable
+``_db=None`` state (its persists would silently no-op). See :class:`SessionRegistry`.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from traceforge.governance.persistence import SystemStore
+    from traceforge.governance.state import SessionState
+
+
+class SessionRegistry:
+    """Owns per-session residency across two scopes: durable and gate-ephemeral.
+
+    Centralizing residency here means no two code paths create divergent state for
+    the same session. But the observation (writer) and gate channels have
+    incompatible needs, so each session keeps a state in whichever of two maps
+    matches the caller's channel — the writer never borrows the gate's state and
+    vice versa:
+
+      * ``_durable_states`` — the observation channel (single writer:
+        ``SessionMonitor``). ``get_or_create`` rehydrates from the durable store,
+        so entries are always DB-backed and their persists durably write.
+      * ``_gate_states`` — the gate channel (``Shield``). ``ensure`` creates
+        thread-safe ephemeral state (``_db=None``, never touching SQLite
+        cross-thread) for enforcement context.
+
+    Because the two channels hold *separate* ``SessionState`` objects per session,
+    their counters can never be mutated concurrently from different threads, and
+    the writer's ``persist_no_commit`` is guaranteed a live connection (the F2
+    durability invariant).
+    """
+
+    def __init__(self, store: "SystemStore") -> None:
+        self._store = store
+        # Observation channel: DB-backed states (single writer: SessionMonitor).
+        self._durable_states: dict[str, "SessionState"] = {}
+        # Gate channel: ephemeral _db=None states (Shield enforcement).
+        self._gate_states: dict[str, "SessionState"] = {}
+
+    # ── Durable / observation residency (the single writer) ──
+
+    def get_or_create(self, session_id: str) -> "SessionState":
+        """Return the durable, DB-backed state, rehydrating from the store on a miss.
+
+        Always yields a state with a live connection, so the writer's
+        ``persist_no_commit`` / ``persist`` durably write. Observation channel only.
+        """
+        from traceforge.governance.state import SessionState
+
+        if session_id not in self._durable_states:
+            self._durable_states[session_id] = SessionState.load_from_db(
+                session_id, self._store.connection, self._store.lock
+            )
+        return self._durable_states[session_id]
+
+    def evict(self, session_id: str) -> "SessionState | None":
+        """Remove and return the durable (observation) state for a session, if any."""
+        return self._durable_states.pop(session_id, None)
+
+    # ── Ephemeral / gate residency (enforcement, thread-safe) ──
+
+    def ensure(self, session_id: str) -> "SessionState":
+        """Return the gate's ephemeral state, creating it (``_db=None``) on a miss.
+
+        Never rehydrates from SQLite, so it is safe to call from any thread.
+        Uses ``setdefault`` so concurrent callers converge on one instance.
+        """
+        from traceforge.governance.state import SessionState
+
+        if session_id not in self._gate_states:
+            self._gate_states.setdefault(session_id, SessionState(session_id=session_id))
+        return self._gate_states[session_id]
+
+    def get_gate(self, session_id: str) -> "SessionState | None":
+        """Return the gate's resident ephemeral state, or None if not resident."""
+        return self._gate_states.get(session_id)
+
+    def evict_gate(self, session_id: str) -> "SessionState | None":
+        """Remove and return the gate's ephemeral state for a session, if any."""
+        return self._gate_states.pop(session_id, None)
+
+    # ── Read-only preview (Scorer) ──
+
+    def preview_state(self, session_id: str) -> "SessionState | None":
+        """Return a base state for a non-persisted preview, or None if unknown.
+
+        Prefers the durable observation state (so a preview reflects accumulated,
+        persisted budgets), falling back to the gate's ephemeral state. The caller
+        clones this via ``clone_detached`` and never mutates or persists it, so
+        returning either residency is side-effect-free.
+        """
+        return self._durable_states.get(session_id) or self._gate_states.get(session_id)

traceforge/governance/results.py

@@ -0,0 +1,183 @@
+"""Governance result types — strongly typed outputs from the enrichment pipeline.
+
+These types are intentionally separated from pipeline.py to break the circular
+import between traceforge.types (EventMetadata) and the governance layer.
+Both types.py and pipeline.py import from here.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from datetime import datetime
+from enum import StrEnum
+from types import MappingProxyType
+
+from traceforge.classify.core import Classification
+from traceforge.classify.risk import RiskAssessment
+from traceforge.governance.drift import DriftAssessment
+from traceforge.governance.state import BudgetSnapshot
+
+
+class RecommendedAction(StrEnum):
+    ALLOW = "allow"
+    WARN = "warn"
+    ESCALATE = "escalate"
+    DENY = "deny"
+    TRANSFORM = "transform"
+
+
+def _empty_parameters() -> Mapping[str, str]:
+    """Default immutable (read-only) parameter mapping."""
+    return MappingProxyType({})
+
+
+@dataclass(frozen=True)
+class TransformSuggestion:
+    """Materialized by Phase 3 from TransformTemplate + event-specific data.
+
+    The transform is *advisory*: it describes what a safe alternative would look
+    like via a ``strategy`` and immutable ``parameters``, plus the resolved
+    ``original_value`` of ``target_field`` (``None`` when the field is absent).
+    traceforge never applies the transform.
+    """
+
+    target_kind: str  # "shell_flag", "shell_arg", "tool_arg", "file_content", "field"
+    path: str  # AST node path (shell) or JSONPath (mcp tool args)
+    original: str
+    replacement: str | None  # None = suggest removal
+    rationale: str
+    confidence: str = "medium"  # "high", "medium", "low"
+    target_field: str | None = None  # dotted path resolved against event data
+    strategy: str = "pattern_replace"  # how to transform: e.g. redact/remove/replace
+    parameters: Mapping[str, str] = field(default_factory=_empty_parameters)
+    original_value: object | None = None  # resolved value of target_field; None if missing
+
+    def __post_init__(self) -> None:
+        # Freeze parameters into a read-only copy so callers cannot mutate them
+        # after construction (frozen=True only guards attribute rebinding).
+        object.__setattr__(self, "parameters", MappingProxyType(dict(self.parameters)))
+
+    def __hash__(self) -> int:
+        # Hashable despite the mapping field. original_value is excluded so the
+        # suggestion stays hashable even when it resolves to a container; this is
+        # consistent with __eq__ (equal objects share all fields -> equal hashes).
+        return hash(
+            (
+                self.target_kind,
+                self.path,
+                self.original,
+                self.replacement,
+                self.rationale,
+                self.confidence,
+                self.target_field,
+                self.strategy,
+                frozenset(self.parameters.items()),
+            )
+        )
+
+
+@dataclass(frozen=True)
+class EscalationContext:
+    """Rich metadata for escalate/deny — full classification context."""
+
+    canonical_id: str
+    classification: "Classification"
+    recommended_action: "RecommendedAction"
+    reason_code: str
+    mitre_techniques: tuple[str, ...]
+    drift: "DriftAssessment | None"
+    budget_snapshot: "BudgetSnapshot | None"
+    pii_taint: bool
+    ifc_violations: int
+    tool_name: str
+    tool_args_summary: str  # Sanitized — no secrets
+    session_id: str
+    timestamp: datetime
+    # --- #24: richer escalation context ---
+    event_id: str = ""  # Triggering event's id (ctx.event.event_id)
+    classification_summary: str = ""  # Concise deterministic human-readable summary
+    risk_factors: tuple[str, ...] = ()  # risk.factors that drove the score
+    session_event_count: int = 0  # snapshot.event_count at escalation time
+    recent_phase_window: tuple[str, ...] = ()  # snapshot.phase_window (recent phases)
+
+
+@dataclass(frozen=True)
+class EvidencePointer:
+    """What triggered this evidence."""
+
+    event_id: str
+    rule_id: str
+    detector: str
+    payload_pointer: str | None = None
+
+
+@dataclass(frozen=True)
+class Evidence:
+    """Emitted for warn/escalate/deny recommendations."""
+
+    canonical_id: str
+    timestamp: datetime
+    session_id: str
+    mechanism: str
+    effect: str | None
+    scope: tuple[str, ...]
+    role: tuple[str, ...]
+    action: tuple[str, ...]
+    capability: tuple[str, ...]
+    structure: tuple[str, ...]
+    source_labels: tuple[str, ...]
+    recommended_action: RecommendedAction
+    risk_score: int
+    risk_factors: tuple[str, ...]
+    mitre_techniques: tuple[str, ...]
+    pointers: tuple[EvidencePointer, ...]
+    escalation: EscalationContext | None = None
+    # --- #25: promote rule provenance + matched-predicate detail to the top level ---
+    rule_id: str = ""  # Matching rule id (also carried on each EvidencePointer)
+    matched_predicates: tuple[str, ...] = ()  # Serialized rule predicates that matched
+
+
+@dataclass(frozen=True)
+class RiskRecommendation:
+    """Full recommendation with canonical identity."""
+
+    recommended_action: RecommendedAction
+    assessment: "RiskAssessment"
+    reason_code: str
+    canonical_id: str
+    message: str | None = None
+    transform: TransformSuggestion | None = None
+
+
+@dataclass(frozen=True)
+class RecommendationResult:
+    """Phase 3 output envelope."""
+
+    recommendation: RiskRecommendation
+    evidence: Evidence | None = None
+
+
+@dataclass(frozen=True)
+class Phase3Result:
+    """Always produced by Phase 3."""
+
+    risk_assessment: "RiskAssessment"
+    recommendation_result: RecommendationResult | None = None
+
+
+@dataclass(frozen=True)
+class SessionMeta:
+    """Full classification output. Attached to event metadata as `governance`.
+
+    For lifecycle events (session_start/end), Phase 2/3 fields are None.
+    canonical_id is accessed via recommendation.canonical_id (no separate field).
+    """
+
+    classification: "Classification | None"
+    risk_assessment: "RiskAssessment | None"
+    recommendation: RiskRecommendation | None = None
+    budget_snapshot: "BudgetSnapshot | None" = None
+    drift: "DriftAssessment | None" = None
+    mcp_alerts: tuple = ()  # tuple[MCPIntegrityAlert, ...]
+    evidence: Evidence | None = None

traceforge/governance/risk_wrapper.py

@@ -0,0 +1,113 @@
+"""Governance risk assessment wrapper around existing assess_risk()."""
+
+from __future__ import annotations
+
+import dataclasses
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from traceforge.classify.config import ClassificationEngine
+    from traceforge.classify.core import Classification
+    from traceforge.classify.risk import RiskAssessment
+    from traceforge.governance.types import CommandAnalysis
+
+
+@dataclass(frozen=True)
+class RiskModifiers:
+    """Governance-specific risk score modifiers (additive bonuses)."""
+
+    phase_drift_bonus: int = 0  # +10 per anomaly up to 20
+    mcp_drift_bonus: int = 0  # +15 per schema change
+    ifc_violations: int = 0  # +10 per violation up to 30
+    integrity_bonus: int = 0  # +10 if integrity_unverified
+    budget_bonus: int = 0  # +5 if budget pressure
+
+
+def assess_governance_risk(
+    enriched_classification: "Classification",
+    command_analysis: "CommandAnalysis | None",
+    risk_modifiers: RiskModifiers,
+    *,
+    engine: "ClassificationEngine",
+    project_root: str | None = None,
+) -> "RiskAssessment":
+    """Compute governance-enriched risk score.
+
+    Wraps existing assess_risk() and adds governance bonuses.
+    """
+    from traceforge.classify.risk import assess_risk
+
+    if command_analysis:
+        base = assess_risk(
+            classification=enriched_classification,
+            command=command_analysis.command or "",
+            engine=engine,
+            binary=command_analysis.binary or "",
+            flags=list(command_analysis.flags) if command_analysis.flags else [],
+            targets=list(command_analysis.targets) if command_analysis.targets else [],
+            pipe_segments=[
+                {"binary": seg.binary, "targets": list(seg.targets)}
+                for seg in command_analysis.pipe_segments
+            ]
+            if command_analysis.pipe_segments
+            else None,
+            project_root=project_root,
+        )
+    else:
+        base = assess_risk(
+            classification=enriched_classification,
+            command="",
+            engine=engine,
+            binary="",
+            flags=[],
+            targets=[],
+            pipe_segments=None,
+            project_root=project_root,
+        )
+
+    # Add governance bonuses (additive, capped at 100)
+    bonus = (
+        risk_modifiers.phase_drift_bonus
+        + risk_modifiers.mcp_drift_bonus
+        + risk_modifiers.integrity_bonus
+        + risk_modifiers.budget_bonus
+    )
+    if risk_modifiers.ifc_violations > 0:
+        bonus += min(risk_modifiers.ifc_violations * 10, 30)
+
+    final_score = min(base.score + bonus, 100)
+
+    # Append governance factors
+    extra_factors: list[str] = []
+    if risk_modifiers.phase_drift_bonus > 0:
+        extra_factors.append(f"phase_drift:+{risk_modifiers.phase_drift_bonus}")
+    if risk_modifiers.mcp_drift_bonus > 0:
+        extra_factors.append(f"mcp_drift:+{risk_modifiers.mcp_drift_bonus}")
+    if risk_modifiers.ifc_violations > 0:
+        extra_factors.append(f"ifc_violations:{risk_modifiers.ifc_violations}")
+    if risk_modifiers.integrity_bonus > 0:
+        extra_factors.append("integrity_unverified")
+    if risk_modifiers.budget_bonus > 0:
+        extra_factors.append("budget_pressure")
+
+    # Compute level from final score
+    level = _score_to_level(final_score)
+
+    return dataclasses.replace(
+        base,
+        score=final_score,
+        level=level,
+        factors=base.factors + tuple(extra_factors),
+    )
+
+
+def _score_to_level(score: int) -> str:
+    """Map 0-100 score to level string."""
+    if score >= 85:
+        return "critical"
+    if score >= 65:
+        return "danger"
+    if score >= 40:
+        return "caution"
+    return "safe"