traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/governance/drift.py

@@ -0,0 +1,196 @@
+"""Phase-aware behavioral drift detection with transition-specific bonuses."""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from traceforge.governance.persistence import SystemStore
+    from traceforge.governance.state import SessionStateSnapshot
+    from traceforge.governance.types import EnrichmentContext
+
+
+# Suspicious transition bonuses (from → to: bonus)
+_TRANSITION_BONUSES: dict[tuple[str, str], int] = {
+    ("testing", "implementation"): 10,  # Verification → destructive implementation
+    ("verification", "implementation"): 10,
+    ("testing", "destructive"): 18,  # Verification → destructive
+    ("verification", "destructive"): 18,
+    ("exploration", "network"): 15,  # Exploration → network write
+    ("exploration", "deployment"): 12,  # Exploration → deployment (skipping impl)
+}
+
+# Oscillation threshold
+_OSCILLATION_THRESHOLD = 5  # >5 transitions in window = +20 bonus
+_OSCILLATION_BONUS = 20
+_WARMUP_EVENTS = 5  # Minimum events before drift detection activates
+
+
+@dataclass(frozen=True)
+class DriftAssessment:
+    """Full drift assessment per the design spec."""
+
+    phase_window: tuple[str, ...]
+    baseline_distribution: tuple[tuple[str, float], ...]  # Sorted pairs (immutable)
+    current_phase: str
+    anomaly_score: float  # 0.0–1.0 deviation from baseline
+    risk_bonus: int  # 0–25 pts added to risk
+    transitions: int  # Phase transitions in window
+    anomaly: bool  # True if risk_bonus > 0
+
+
+class DriftDetector:
+    """Stateless drift detector. Phase window from session state, baseline from store."""
+
+    def __init__(self, store: "SystemStore", window_size: int = 20, threshold: float = 0.3) -> None:
+        self._store = store
+        self._window_size = window_size
+        self._threshold = threshold
+
+    def detect(
+        self,
+        ctx: "EnrichmentContext",
+        state_snapshot: "SessionStateSnapshot",
+        cap: set[str],
+    ) -> DriftAssessment | None:
+        """Full drift detection: divergence + transition bonuses + oscillation."""
+        phase_window = state_snapshot.phase_window
+        if len(phase_window) < _WARMUP_EVENTS:
+            return None
+
+        current_phase = phase_window[-1] if phase_window else "unknown"
+
+        # Count transitions in window
+        transitions = self._count_transitions(phase_window)
+
+        # Compute distribution divergence
+        current_dist = self._compute_distribution(phase_window)
+        baseline = self._get_baseline(ctx)
+
+        divergence = 0.0
+        baseline_tuples: tuple[tuple[str, float], ...] = ()
+
+        if baseline:
+            baseline_dist = baseline["phase_counts"]
+            total = baseline["total_events"]
+            if total > 0:
+                norm_baseline = {k: v / total for k, v in baseline_dist.items()}
+                baseline_tuples = tuple(sorted(norm_baseline.items()))
+                divergence = self._js_divergence(current_dist, norm_baseline)
+
+        # Compute risk bonus from multiple sources
+        risk_bonus = 0
+
+        # 1. Transition-specific bonuses
+        if len(phase_window) >= 2:
+            prev_phase = phase_window[-2]
+            transition_key = (prev_phase, current_phase)
+            risk_bonus += _TRANSITION_BONUSES.get(transition_key, 0)
+
+        # 2. Oscillation bonus
+        if transitions > _OSCILLATION_THRESHOLD:
+            risk_bonus += _OSCILLATION_BONUS
+
+        # 3. Divergence-based bonus (scale anomaly_score to 0-15 pts)
+        if divergence > self._threshold:
+            risk_bonus += min(int(divergence * 25), 15)
+
+        # Cap total at 25
+        risk_bonus = min(risk_bonus, 25)
+
+        anomaly = risk_bonus > 0
+
+        if anomaly:
+            # phase_anomaly tracked via structure only (not capability)
+            # to avoid polluting canonical hash with session-contextual state
+            pass
+
+        return DriftAssessment(
+            phase_window=phase_window,
+            baseline_distribution=baseline_tuples,
+            current_phase=current_phase,
+            anomaly_score=divergence,
+            risk_bonus=risk_bonus,
+            transitions=transitions,
+            anomaly=anomaly,
+        )
+
+    def check_drift(
+        self,
+        phase_window: tuple[str, ...],
+        current_phase: str,
+        baseline: tuple[tuple[str, float], ...] | None,
+    ) -> DriftAssessment | None:
+        """Pure function version — receives pre-loaded baseline."""
+        if len(phase_window) < _WARMUP_EVENTS:
+            return None
+
+        transitions = self._count_transitions(phase_window)
+        current_dist = self._compute_distribution(phase_window)
+
+        divergence = 0.0
+        if baseline:
+            norm_baseline = dict(baseline)
+            divergence = self._js_divergence(current_dist, norm_baseline)
+
+        risk_bonus = 0
+        if len(phase_window) >= 2:
+            prev_phase = phase_window[-2]
+            risk_bonus += _TRANSITION_BONUSES.get((prev_phase, current_phase), 0)
+        if transitions > _OSCILLATION_THRESHOLD:
+            risk_bonus += _OSCILLATION_BONUS
+        if divergence > self._threshold:
+            risk_bonus += min(int(divergence * 25), 15)
+        risk_bonus = min(risk_bonus, 25)
+
+        return DriftAssessment(
+            phase_window=phase_window,
+            baseline_distribution=baseline or (),
+            current_phase=current_phase,
+            anomaly_score=divergence,
+            risk_bonus=risk_bonus,
+            transitions=transitions,
+            anomaly=risk_bonus > 0,
+        )
+
+    def _get_baseline(self, ctx: "EnrichmentContext") -> dict | None:
+        # Use pre-loaded baseline from context when available
+        if ctx.drift_baseline:
+            return {
+                "phase_counts": dict(ctx.drift_baseline),
+                "total_events": sum(v for _, v in ctx.drift_baseline),
+            }
+        repo = ctx.project_root or "unknown"
+        return self._store.get_drift_baseline("unknown", repo)
+
+    def _count_transitions(self, phases: tuple[str, ...]) -> int:
+        """Count phase transitions (consecutive different phases) in window."""
+        count = 0
+        for i in range(1, len(phases)):
+            if phases[i] != phases[i - 1]:
+                count += 1
+        return count
+
+    def _compute_distribution(self, phases: tuple[str, ...] | list[str]) -> dict[str, float]:
+        counts: dict[str, int] = {}
+        for p in phases:
+            counts[p] = counts.get(p, 0) + 1
+        total = len(phases)
+        return {k: v / total for k, v in counts.items()}
+
+    def _js_divergence(self, p: dict[str, float], q: dict[str, float]) -> float:
+        """Jensen-Shannon divergence (symmetric, bounded [0, ln(2)])."""
+        all_keys = set(p.keys()) | set(q.keys())
+        eps = 1e-10
+        divergence = 0.0
+        for k in all_keys:
+            pk = p.get(k, eps)
+            qk = q.get(k, eps)
+            m = (pk + qk) / 2
+            if pk > eps:
+                divergence += pk * math.log(pk / m)
+            if qk > eps:
+                divergence += qk * math.log(qk / m)
+        return divergence / 2

traceforge/governance/emitter.py

@@ -0,0 +1,234 @@
+"""Async emission actor for governance-enriched events, with backpressure.
+
+The observer enriches **synchronously** (single-writer governance: ``observe_event``
+runs exactly once per event and advances the tool-call budget) and hands an
+already-scored ``(event, meta)`` pair to this actor via :meth:`EnrichedEmitter.submit`.
+The actor owns *only* audit emission and backpressure — it never enriches, so it
+can never re-run ``observe_event`` and double-count the budget.
+
+A bounded :class:`asyncio.Queue` decouples the latency-sensitive enforcement path
+(the observer hook returns ``SessionMeta`` to the host immediately) from
+potentially slow sinks. When the queue is full the **oldest audit record** is
+dropped — never an enforcement decision, which already took effect synchronously
+in the observer — the drop is counted durably via an injected ``record_drop``
+callback, and a **coalesced** :class:`ContextGapEvent` is emitted downstream so
+consumers see an explicit gap marker instead of silent loss.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import uuid
+from datetime import datetime, timezone
+from typing import TYPE_CHECKING, Callable, Iterable
+
+from traceforge.governance.envelope import ContextGapEvent, EnrichedEvent
+from traceforge.governance.results import SessionMeta
+
+if TYPE_CHECKING:
+    import traceforge.types
+    from traceforge.sinks.base import StorageSink
+
+logger = logging.getLogger(__name__)
+
+# Synthetic gap markers bypass enrichment, so they carry an empty SessionMeta
+# (serializes to ``"_governance": {}``). Reused — SessionMeta is immutable data.
+_EMPTY_META = SessionMeta(classification=None, risk_assessment=None)
+
+DEFAULT_CAPACITY = 1024
+
+
+class _GapAccumulator:
+    """Coalesces consecutive dropped events for one session into one marker.
+
+    The durable drop *counter* is precise (``record_drop`` fires once per dropped
+    event); the emitted *marker* is coalesced so a burst of drops surfaces as a
+    single :class:`ContextGapEvent` spanning ``first``..``last`` sequence.
+    """
+
+    __slots__ = ("session_id", "count", "first_sequence", "last_sequence", "gap_ordinal")
+
+    def __init__(self, session_id: str) -> None:
+        self.session_id = session_id
+        self.count = 0
+        self.first_sequence: int | None = None
+        self.last_sequence: int | None = None
+        self.gap_ordinal = 0
+
+    def add(self, sequence: int | None, gap_ordinal: int) -> None:
+        self.count += 1
+        self.gap_ordinal = gap_ordinal
+        if sequence is not None:
+            if self.first_sequence is None:
+                self.first_sequence = sequence
+            self.last_sequence = sequence
+
+    def to_event(self) -> ContextGapEvent:
+        key = ContextGapEvent.compute_source_event_key(
+            self.session_id,
+            self.first_sequence,
+            self.last_sequence,
+            self.gap_ordinal,
+        )
+        return ContextGapEvent(
+            id=f"gap-{uuid.uuid4().hex[:12]}",
+            session_id=self.session_id,
+            timestamp=datetime.now(timezone.utc),
+            source_event_key=key,
+            dropped_count=self.count,
+            first_dropped_sequence=self.first_sequence,
+            last_dropped_sequence=self.last_sequence,
+            gap_ordinal=self.gap_ordinal,
+        )
+
+
+class EnrichedEmitter:
+    """Bounded async actor: wraps ``(event, meta)`` in :class:`EnrichedEvent` and
+    fans it out to sinks via ``on_enriched_event``, dropping the oldest audit
+    record under backpressure.
+
+    Runs on a single event loop; :meth:`submit` is synchronous and non-blocking
+    (safe to call from inside the observer's async hooks). Because everything runs
+    single-threaded on the loop, ``submit`` executes atomically with respect to the
+    drain task — no locks are needed.
+    """
+
+    def __init__(
+        self,
+        sinks: "Iterable[StorageSink]",
+        *,
+        capacity: int = DEFAULT_CAPACITY,
+        record_drop: "Callable[[str, int], None] | None" = None,
+    ) -> None:
+        if capacity < 1:
+            raise ValueError("capacity must be >= 1")
+        self._sinks: list[StorageSink] = list(sinks)
+        self._capacity = capacity
+        self._record_drop = record_drop
+        self._queue: asyncio.Queue = asyncio.Queue(maxsize=capacity)
+        # Per-session coalescing of dropped events, flushed just before that
+        # session's next surviving event (and any trailing remainder at aclose).
+        self._pending_gaps: dict[str, _GapAccumulator] = {}
+        self._gap_ordinals: dict[str, int] = {}
+        self._drain_task: asyncio.Task | None = None
+        self._closed = False
+
+    @property
+    def capacity(self) -> int:
+        return self._capacity
+
+    async def start(self) -> None:
+        """Begin draining. Idempotent; must be called inside the event loop that
+        will own emission."""
+        if self._drain_task is None and not self._closed:
+            self._drain_task = asyncio.create_task(self._drain_loop())
+
+    def submit(self, event: "traceforge.types.SessionEvent", meta: SessionMeta) -> None:
+        """Enqueue an already-enriched ``(event, meta)`` for emission.
+
+        Synchronous and non-blocking. On a full queue, drop the **oldest**
+        surviving audit record (enforcement already happened in the observer, so
+        only audit emission is lost), record it durably, coalesce it into the
+        session's pending gap marker, then enqueue the new item.
+        """
+        item = (event, meta)
+        try:
+            self._queue.put_nowait(item)
+            return
+        except asyncio.QueueFull:
+            pass
+        try:
+            dropped_event, _ = self._queue.get_nowait()
+            # Balance the unfinished-task count for the put() that enqueued the
+            # dropped item, so aclose()'s queue.join() cannot hang.
+            self._queue.task_done()
+        except asyncio.QueueEmpty:
+            dropped_event = None
+        if dropped_event is not None:
+            self._record_dropped(dropped_event)
+        try:
+            self._queue.put_nowait(item)
+        except asyncio.QueueFull:
+            # We just freed a slot, so this should not happen; if it somehow does,
+            # count the *new* item as dropped rather than block the caller.
+            self._record_dropped(event)
+
+    def _record_dropped(self, event: "traceforge.types.SessionEvent") -> None:
+        sid = getattr(event, "session_id", "") or ""
+        sequence = None
+        metadata = getattr(event, "metadata", None)
+        if metadata is not None:
+            sequence = getattr(metadata, "sequence", None)
+        ordinal = self._gap_ordinals.get(sid, 0) + 1
+        self._gap_ordinals[sid] = ordinal
+        acc = self._pending_gaps.get(sid)
+        if acc is None:
+            acc = _GapAccumulator(sid)
+            self._pending_gaps[sid] = acc
+        acc.add(sequence, ordinal)
+        if self._record_drop is not None:
+            try:
+                self._record_drop(sid, 1)
+            except Exception as exc:  # never let a persistence hiccup break submit
+                logger.error("record_drop failed for session %s: %s", sid, exc)
+
+    async def _drain_loop(self) -> None:
+        while True:
+            event, meta = await self._queue.get()
+            try:
+                await self._emit(event, meta)
+            finally:
+                self._queue.task_done()
+
+    async def _emit(self, event: "traceforge.types.SessionEvent", meta: SessionMeta) -> None:
+        # Flush this session's coalesced gap (if any) *before* the surviving event,
+        # so downstream sees the gap marker in order ahead of the next real record.
+        sid = getattr(event, "session_id", "") or ""
+        acc = self._pending_gaps.pop(sid, None)
+        if acc is not None:
+            await self._fanout_enriched(EnrichedEvent(event=acc.to_event(), governance=_EMPTY_META))
+        await self._fanout_enriched(EnrichedEvent(event=event, governance=meta))
+
+    async def _fanout_enriched(self, enriched: EnrichedEvent) -> None:
+        results = await asyncio.gather(
+            *(sink.on_enriched_event(enriched) for sink in self._sinks),
+            return_exceptions=True,
+        )
+        self._log_sink_errors(results, "on_enriched_event")
+
+    def _log_sink_errors(self, results: list, op: str) -> None:
+        for sink, result in zip(self._sinks, results):
+            if isinstance(result, BaseException):
+                logger.error(
+                    "Sink %s failed during %s: %s",
+                    type(sink).__name__,
+                    op,
+                    result,
+                    exc_info=(type(result), result, result.__traceback__),
+                )
+
+    async def aclose(self) -> None:
+        """Drain everything already submitted, flush trailing gap markers, then
+        flush sinks. Idempotent."""
+        if self._closed:
+            return
+        self._closed = True
+        if self._drain_task is not None:
+            # Wait for all queued items to be emitted, then stop the loop.
+            await self._queue.join()
+            self._drain_task.cancel()
+            try:
+                await self._drain_task
+            except asyncio.CancelledError:
+                pass
+            self._drain_task = None
+        # Any coalesced gaps whose surviving event never arrived still deserve a
+        # downstream marker.
+        for sid in list(self._pending_gaps.keys()):
+            acc = self._pending_gaps.pop(sid)
+            await self._fanout_enriched(EnrichedEvent(event=acc.to_event(), governance=_EMPTY_META))
+        results = await asyncio.gather(
+            *(sink.flush() for sink in self._sinks), return_exceptions=True
+        )
+        self._log_sink_errors(results, "flush")

traceforge/governance/envelope.py

@@ -0,0 +1,138 @@
+"""Enriched event envelope, ContextGapEvent, and sink emission types."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import datetime
+from typing import TYPE_CHECKING, Literal
+
+if TYPE_CHECKING:
+    import traceforge.types
+
+    from traceforge.governance.pipeline import SessionMeta
+    from traceforge.governance.types import SessionEvent
+
+
+@dataclass(frozen=True)
+class ContextGapEvent:
+    """Synthetic marker emitted when events are dropped due to backpressure.
+
+    Does NOT flow through full enrichment pipeline — bypasses Phase 2/3.
+    Serialized directly to sinks as-is.
+    """
+
+    id: str
+    session_id: str
+    timestamp: datetime
+    source_event_key: str
+    kind: Literal["context_gap"] = "context_gap"
+    dropped_count: int = 0
+    first_dropped_sequence: int | None = None
+    last_dropped_sequence: int | None = None
+    gap_ordinal: int = 0
+    reason: str = "backpressure"
+
+    @staticmethod
+    def compute_source_event_key(
+        session_id: str,
+        first_sequence: int | None,
+        last_sequence: int | None,
+        gap_ordinal: int,
+    ) -> str:
+        """Deterministic key derivation for gap events."""
+        if first_sequence is not None and last_sequence is not None:
+            return f"gap:{session_id}:{first_sequence}:{last_sequence}"
+        return f"gap:{session_id}:ord:{gap_ordinal}"
+
+
+@dataclass(frozen=True)
+class EnrichedEvent:
+    """Immutable envelope for sink emission. Event is unmodified; governance is attached alongside.
+
+    Sinks serialize this as {"event": {...}, "_governance": {...}} or equivalent per sink type.
+    """
+
+    event: "traceforge.types.SessionEvent | SessionEvent | ContextGapEvent"
+    governance: "SessionMeta"
+
+    def to_dict(self) -> dict:
+        """Serialize for sink emission."""
+        event_dict: dict
+        if isinstance(self.event, ContextGapEvent):
+            event_dict = {
+                "kind": self.event.kind,
+                "id": self.event.id,
+                "session_id": self.event.session_id,
+                "timestamp": self.event.timestamp.isoformat(),
+                "dropped_count": self.event.dropped_count,
+                "first_dropped_sequence": self.event.first_dropped_sequence,
+                "last_dropped_sequence": self.event.last_dropped_sequence,
+                "gap_ordinal": self.event.gap_ordinal,
+                "reason": self.event.reason,
+            }
+        elif hasattr(self.event, "model_dump"):
+            # Live traceforge.types.SessionEvent (pydantic). Mirror JsonlSink's
+            # event body, but lift governance out of metadata into the top-level
+            # ``_governance`` slot so the event stays the raw structural record
+            # and governance is not duplicated inside it.
+            live = self.event
+            metadata_dict = (
+                live.metadata.model_dump(exclude_none=True) if live.metadata is not None else None
+            )
+            if isinstance(metadata_dict, dict):
+                metadata_dict.pop("governance", None)
+            event_dict = {
+                "id": live.id,
+                "kind": live.kind,
+                "session_id": live.session_id,
+                "timestamp": live.timestamp.isoformat(),
+                "payload": live.payload,
+                "metadata": metadata_dict,
+            }
+        else:
+            event_dict = {
+                "event_id": self.event.event_id,
+                "session_id": self.event.session_id,
+                "timestamp": self.event.timestamp.isoformat(),
+                "source_event_key": self.event.source_event_key,
+            }
+
+        governance_dict = {}
+        if self.governance.classification is not None:
+            governance_dict["classification"] = self.governance.classification.to_dict()
+        if self.governance.risk_assessment is not None:
+            ra = self.governance.risk_assessment
+            governance_dict["risk_assessment"] = {
+                "score": ra.score,
+                "level": ra.level,
+                "confidence": ra.confidence,
+                "factors": list(ra.factors),
+                "mitre": list(ra.mitre),
+            }
+        if self.governance.recommendation is not None:
+            rec = self.governance.recommendation
+            governance_dict["recommendation"] = {
+                "action": rec.recommended_action.value,
+                "reason_code": rec.reason_code,
+                "canonical_id": rec.canonical_id,
+            }
+        if self.governance.evidence is not None:
+            ev = self.governance.evidence
+            governance_dict["evidence"] = {
+                "canonical_id": ev.canonical_id,
+                "recommended_action": ev.recommended_action.value,
+                "risk_score": ev.risk_score,
+                "mechanism": ev.mechanism,
+                "effect": ev.effect,
+                "scope": list(ev.scope),
+                "capability": list(ev.capability),
+                "structure": list(ev.structure),
+            }
+        if self.governance.budget_snapshot is not None:
+            bs = self.governance.budget_snapshot
+            governance_dict["budget"] = {
+                "total_tool_calls": bs.total_tool_calls,
+                "pressure": bs.pressure,
+            }
+
+        return {"event": event_dict, "_governance": governance_dict}