traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/gate/registry.py

@@ -0,0 +1,108 @@
+"""Gate endpoint registry — maps session_id → socket path in system.db.
+
+The gate_endpoints table is managed by Alembic as part of the unified
+system schema. This module provides convenience functions that operate
+on a standalone connection for lightweight lookups (the gate client doesn't
+always have access to a full SystemStore instance).
+"""
+
+from __future__ import annotations
+
+import os
+import sqlite3
+from pathlib import Path
+
+
+def _gates_dir() -> Path:
+    """Return ~/.traceforge/gates/, creating if needed."""
+    d = Path.home() / ".traceforge" / "gates"
+    d.mkdir(parents=True, exist_ok=True)
+    return d
+
+
+def _system_db_path() -> str:
+    return str(Path.home() / ".traceforge" / "system.db")
+
+
+def register_session(session_id: str, sock_path: str, *, db_path: str | None = None) -> None:
+    """Register a session_id → sock_path mapping."""
+    db = db_path or _system_db_path()
+    Path(db).parent.mkdir(parents=True, exist_ok=True)
+    conn = sqlite3.connect(db)
+    try:
+        conn.execute(
+            "INSERT OR REPLACE INTO gate_endpoints (session_id, sock_path, pid) VALUES (?, ?, ?)",
+            (session_id, sock_path, os.getpid()),
+        )
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def lookup_session(session_id: str, *, db_path: str | None = None) -> str | None:
+    """Look up the socket path for a session_id. Returns None if not found or stale."""
+    db = db_path or _system_db_path()
+    if not Path(db).exists():
+        return None
+    conn = sqlite3.connect(db)
+    try:
+        row = conn.execute(
+            "SELECT sock_path, pid FROM gate_endpoints WHERE session_id = ?",
+            (session_id,),
+        ).fetchone()
+        if row is None:
+            return None
+        sock_path, pid = row
+        # Verify PID is alive
+        if not _pid_alive(pid):
+            conn.execute("DELETE FROM gate_endpoints WHERE session_id = ?", (session_id,))
+            conn.commit()
+            return None
+        return sock_path
+    finally:
+        conn.close()
+
+
+def unregister_session(session_id: str, *, db_path: str | None = None) -> None:
+    """Remove a session_id from the registry."""
+    db = db_path or _system_db_path()
+    if not Path(db).exists():
+        return
+    conn = sqlite3.connect(db)
+    try:
+        conn.execute("DELETE FROM gate_endpoints WHERE session_id = ?", (session_id,))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def unregister_pid(*, db_path: str | None = None) -> None:
+    """Remove all entries for the current PID (cleanup on exit)."""
+    db = db_path or _system_db_path()
+    if not Path(db).exists():
+        return
+    conn = sqlite3.connect(db)
+    try:
+        conn.execute("DELETE FROM gate_endpoints WHERE pid = ?", (os.getpid(),))
+        conn.commit()
+    finally:
+        conn.close()
+
+
+def _pid_alive(pid: int) -> bool:
+    """Check if a PID is still running."""
+    if os.name == "nt":
+        import ctypes
+
+        kernel32 = ctypes.windll.kernel32
+        handle = kernel32.OpenProcess(0x1000, False, pid)  # PROCESS_QUERY_LIMITED_INFORMATION
+        if handle:
+            kernel32.CloseHandle(handle)
+            return True
+        return False
+    else:
+        try:
+            os.kill(pid, 0)
+            return True
+        except OSError:
+            return False

traceforge/gate/server.py

@@ -0,0 +1,174 @@
+"""Gate IPC server — listens on a unix socket (or named pipe on Windows) for gate requests.
+
+The server runs in a background thread inside the Pipeline process. When a gate
+request arrives (from `traceforge gate --stdin`), it:
+  1. Deserializes the event JSON
+  2. Calls pipeline._score_and_gate_preflight(payload)
+  3. Returns the Verdict as JSON
+
+Uses the same policy chain as all gate_* methods for consistent behavior.
+"""
+
+from __future__ import annotations
+
+import atexit
+import json
+import os
+import socket
+import struct
+import sys
+import threading
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from traceforge.governance.pipeline import GovernancePipeline
+
+
+class GateServer:
+    """IPC server for cross-process gating.
+
+    Used by CLI-based frameworks (Claude Code, Copilot CLI, Codex CLI, etc.)
+    that can't inject Python hooks but can shell out to `traceforge gate --stdin`.
+
+    The server uses the pipeline's policy chain (same as gate_* methods) rather
+    than a standalone callback. This ensures consistent behavior across all paths.
+    """
+
+    def __init__(
+        self,
+        pipeline: "GovernancePipeline",
+        sock_path: str | None = None,
+    ) -> None:
+        self._pipeline = pipeline
+        self._sock_path = sock_path or self._default_sock_path()
+        self._server: socket.socket | None = None
+        self._thread: threading.Thread | None = None
+        self._running = False
+
+    @staticmethod
+    def _default_sock_path() -> str:
+        gates_dir = Path.home() / ".traceforge" / "gates"
+        gates_dir.mkdir(parents=True, exist_ok=True)
+        return str(gates_dir / f"{os.getpid()}.sock")
+
+    @property
+    def sock_path(self) -> str:
+        return self._sock_path
+
+    def start(self) -> None:
+        """Start the IPC server in a background daemon thread."""
+        if self._running:
+            return
+
+        # Clean up stale socket file
+        if os.path.exists(self._sock_path):
+            os.unlink(self._sock_path)
+
+        if sys.platform == "win32":
+            # Windows: use TCP on localhost with a random port
+            self._server = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+            self._server.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+            self._server.bind(("127.0.0.1", 0))
+            # Store actual port in sock_path for registry
+            _, port = self._server.getsockname()
+            self._sock_path = f"tcp://127.0.0.1:{port}"
+        else:
+            self._server = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
+            self._server.bind(self._sock_path)
+
+        self._server.listen(16)
+        self._server.settimeout(1.0)  # allow periodic check of _running
+        self._running = True
+
+        self._thread = threading.Thread(
+            target=self._serve_loop, daemon=True, name="traceforge-gate"
+        )
+        self._thread.start()
+
+        atexit.register(self.stop)
+
+    def stop(self) -> None:
+        """Stop the IPC server and clean up."""
+        self._running = False
+        if self._server:
+            try:
+                self._server.close()
+            except OSError:
+                pass
+        if self._thread and self._thread.is_alive():
+            self._thread.join(timeout=2.0)
+        # Clean up socket file
+        if not self._sock_path.startswith("tcp://") and os.path.exists(self._sock_path):
+            try:
+                os.unlink(self._sock_path)
+            except OSError:
+                pass
+
+    def register_session(self, session_id: str) -> None:
+        """Register a session_id → this server's socket in the registry."""
+        from traceforge.gate.registry import register_session
+
+        register_session(session_id, self._sock_path)
+
+    def _serve_loop(self) -> None:
+        while self._running:
+            try:
+                conn, _ = self._server.accept()
+            except socket.timeout:
+                continue
+            except OSError:
+                break
+            threading.Thread(target=self._handle_conn, args=(conn,), daemon=True).start()
+
+    def _handle_conn(self, conn: socket.socket) -> None:
+        """Handle a single gate request."""
+        try:
+            conn.settimeout(10.0)
+            data = self._recv_all(conn)
+            if not data:
+                return
+
+            request = json.loads(data)
+            payload = request.get("payload", request)
+            response = self._process_gate_request(payload)
+            resp_bytes = json.dumps(response).encode("utf-8")
+            conn.sendall(struct.pack("!I", len(resp_bytes)) + resp_bytes)
+        except Exception:
+            try:
+                err = json.dumps({"decision": "deny", "reason": "internal error"}).encode()
+                conn.sendall(struct.pack("!I", len(err)) + err)
+            except OSError:
+                pass
+        finally:
+            conn.close()
+
+    def _process_gate_request(self, payload: dict) -> dict:
+        """Score and gate a tool call using the pipeline's policy chain."""
+        trace, verdict = self._pipeline._score_and_gate_preflight(payload)
+        return {
+            "decision": verdict.decision.value,
+            "reason": verdict.reason,
+            "score": trace.risk_score,
+            "level": trace.risk_band,
+        }
+
+    @staticmethod
+    def _recv_all(conn: socket.socket) -> bytes:
+        """Read a length-prefixed message (4-byte big-endian length + payload)."""
+        header = b""
+        while len(header) < 4:
+            chunk = conn.recv(4 - len(header))
+            if not chunk:
+                return b""
+            header += chunk
+        length = struct.unpack("!I", header)[0]
+        if length > 10 * 1024 * 1024:  # 10MB sanity limit
+            return b""
+        data = b""
+        while len(data) < length:
+            chunk = conn.recv(min(length - len(data), 65536))
+            if not chunk:
+                return b""
+            data += chunk
+        return data

traceforge/gates/__init__.py

@@ -0,0 +1,8 @@
+"""Built-in gate implementations for common policy patterns.
+
+Usage:
+    from traceforge.gates.pii import pii_postflight_gate, PiiGateConfig
+    from traceforge.sdk import GatePolicy
+
+    policy = GatePolicy().postflight(pii_postflight_gate())
+"""

traceforge/gates/pii.py

@@ -0,0 +1,352 @@
+"""PII detection and redaction postflight gate — zero external dependencies.
+
+Regex patterns extracted from Microsoft Presidio recognizers + traceforge additions
+(API keys, secrets). Patterns live in pii_patterns.yaml alongside this module.
+
+No NLP model, no spaCy, no Presidio dependency. Pure regex + context boost + checksum
+validation (Luhn for credit cards, mod-97 for IBAN).
+
+Usage:
+    from traceforge.gates.pii import pii_postflight_gate, PiiGateConfig
+    from traceforge.sdk import GatePolicy
+
+    policy = GatePolicy().postflight(pii_postflight_gate())
+
+    # Custom: stricter threshold, only specific entities
+    config = PiiGateConfig(
+        score_threshold=0.7,
+        entities=("US_SSN", "CREDIT_CARD", "API_KEY"),
+        allow_list=("myservice.internal",),
+    )
+    policy = GatePolicy().postflight(pii_postflight_gate(config))
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass
+from pathlib import Path
+from typing import TYPE_CHECKING
+
+import yaml
+
+if TYPE_CHECKING:
+    from traceforge.sdk.gate_types import GateContext, PostflightVerdict, ToolCallResult
+    from traceforge.sdk.verdict import PostflightGate
+
+
+# ─── Configuration ────────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True)
+class PiiGateConfig:
+    """Configuration for the PII postflight gate.
+
+    Attributes:
+        score_threshold: Minimum final confidence (0.0–1.0) for a detection to
+            count as a finding. Default 0.5.
+        entities: Which entity types to scan for. None = all loaded from YAML.
+        critical_entities: Entity types that trigger SUPPRESS. None = use YAML 'critical' field.
+        allow_list: Additional values to never flag (merged with YAML default_allow_list).
+        suppress_on_critical: If True, critical PII → SUPPRESS. If False, always REDACT.
+    """
+
+    score_threshold: float = 0.5
+    entities: tuple[str, ...] | None = None
+    critical_entities: tuple[str, ...] | None = None
+    allow_list: tuple[str, ...] = ()
+    suppress_on_critical: bool = True
+
+
+# ─── Pattern Data ─────────────────────────────────────────────────────────────
+
+
+@dataclass(frozen=True, slots=True)
+class _CompiledPattern:
+    name: str
+    regex: re.Pattern[str]
+    base_score: float
+
+
+@dataclass(frozen=True, slots=True)
+class _EntityDef:
+    entity_type: str
+    critical: bool
+    context_words: frozenset[str]
+    patterns: tuple[_CompiledPattern, ...]
+    validator: str | None
+
+
+@dataclass(frozen=True, slots=True)
+class PiiMatch:
+    """A single PII detection result."""
+
+    entity_type: str
+    start: int
+    end: int
+    score: float
+    pattern_name: str
+    text: str
+
+
+# ─── YAML Loading (module-level singleton) ────────────────────────────────────
+
+_DATA: dict | None = None
+_ENTITIES: tuple[_EntityDef, ...] = ()
+_ALLOW_SET: frozenset[str] = frozenset()
+_CONTEXT_BOOST: float = 0.25
+_CONTEXT_WINDOW: int = 50
+
+
+def _load_patterns() -> None:
+    """Load and compile patterns from pii_patterns.yaml. Called once."""
+    global _DATA, _ENTITIES, _ALLOW_SET, _CONTEXT_BOOST, _CONTEXT_WINDOW
+
+    yaml_path = Path(__file__).parent / "pii_patterns.yaml"
+    with open(yaml_path, "r", encoding="utf-8") as f:
+        _DATA = yaml.safe_load(f)
+
+    _CONTEXT_BOOST = _DATA.get("context_boost", 0.25)
+    _CONTEXT_WINDOW = _DATA.get("context_window", 50)
+    _ALLOW_SET = frozenset(v.lower() for v in _DATA.get("default_allow_list", []))
+
+    entities: list[_EntityDef] = []
+    for ent in _DATA.get("entities", []):
+        compiled: list[_CompiledPattern] = []
+        for p in ent.get("patterns", []):
+            compiled.append(
+                _CompiledPattern(
+                    name=p["name"],
+                    regex=re.compile(p["regex"], re.IGNORECASE),
+                    base_score=p["score"],
+                )
+            )
+        entities.append(
+            _EntityDef(
+                entity_type=ent["entity_type"],
+                critical=ent.get("critical", False),
+                context_words=frozenset(w.lower() for w in ent.get("context", [])),
+                patterns=tuple(compiled),
+                validator=ent.get("validator"),
+            )
+        )
+    _ENTITIES = tuple(entities)
+
+
+def _ensure_loaded() -> None:
+    if _DATA is None:
+        _load_patterns()
+
+
+# ─── Validators ───────────────────────────────────────────────────────────────
+
+
+def _luhn_check(digits: str) -> bool:
+    """Luhn algorithm for credit card validation."""
+    cleaned = re.sub(r"[\s\-]", "", digits)
+    if not cleaned.isdigit() or len(cleaned) < 12:
+        return False
+    total = 0
+    for i, ch in enumerate(reversed(cleaned)):
+        n = int(ch)
+        if i % 2 == 1:
+            n *= 2
+            if n > 9:
+                n -= 9
+        total += n
+    return total % 10 == 0
+
+
+def _iban_check(value: str) -> bool:
+    """IBAN mod-97 validation (ISO 7064)."""
+    cleaned = re.sub(r"\s", "", value).upper()
+    if len(cleaned) < 15 or not cleaned[:2].isalpha() or not cleaned[2:4].isdigit():
+        return False
+    # Move country + check digits to end
+    rearranged = cleaned[4:] + cleaned[:4]
+    # Convert letters to digits (A=10, B=11, ..., Z=35)
+    numeric = ""
+    for ch in rearranged:
+        if ch.isdigit():
+            numeric += ch
+        else:
+            numeric += str(ord(ch) - 55)
+    try:
+        return int(numeric) % 97 == 1
+    except (ValueError, OverflowError):
+        return False
+
+
+_VALIDATORS: dict[str, callable] = {
+    "luhn": _luhn_check,
+    "iban": _iban_check,
+}
+
+
+# ─── Detection Engine ─────────────────────────────────────────────────────────
+
+
+def scan_text(
+    text: str,
+    *,
+    entities: tuple[str, ...] | None = None,
+    score_threshold: float = 0.5,
+    allow_list: frozenset[str] = frozenset(),
+) -> list[PiiMatch]:
+    """Scan text for PII. Returns matches above score_threshold.
+
+    This is the core detection function — usable standalone or via the gate.
+    """
+    _ensure_loaded()
+
+    merged_allow = _ALLOW_SET | frozenset(v.lower() for v in allow_list)
+    text_lower = text.lower()
+    matches: list[PiiMatch] = []
+
+    for entity_def in _ENTITIES:
+        if entities and entity_def.entity_type not in entities:
+            continue
+
+        for pat in entity_def.patterns:
+            for m in pat.regex.finditer(text):
+                matched_text = m.group()
+
+                # Allow list check
+                if matched_text.lower() in merged_allow:
+                    continue
+
+                # Base score
+                score = pat.base_score
+
+                # Context boost: check surrounding text for context words
+                start_ctx = max(0, m.start() - _CONTEXT_WINDOW)
+                end_ctx = min(len(text), m.end() + _CONTEXT_WINDOW)
+                window = text_lower[start_ctx:end_ctx]
+                if entity_def.context_words:
+                    for cw in entity_def.context_words:
+                        if cw in window:
+                            score = min(1.0, score + _CONTEXT_BOOST)
+                            break  # One context match is enough
+
+                # Validator (checksum) → promotes to 1.0 or demotes to 0
+                if entity_def.validator and entity_def.validator in _VALIDATORS:
+                    validator_fn = _VALIDATORS[entity_def.validator]
+                    if validator_fn(matched_text):
+                        score = 1.0
+                    else:
+                        score = max(0.0, score - 0.3)
+
+                if score >= score_threshold:
+                    matches.append(
+                        PiiMatch(
+                            entity_type=entity_def.entity_type,
+                            start=m.start(),
+                            end=m.end(),
+                            score=score,
+                            pattern_name=pat.name,
+                            text=matched_text,
+                        )
+                    )
+
+    # Deduplicate overlapping spans — higher score wins
+    matches.sort(key=lambda x: (-x.score, x.start))
+    deduped: list[PiiMatch] = []
+    taken_ranges: list[tuple[int, int]] = []
+    for match in matches:
+        overlaps = any(match.start < end and match.end > start for start, end in taken_ranges)
+        if not overlaps:
+            deduped.append(match)
+            taken_ranges.append((match.start, match.end))
+
+    return deduped
+
+
+# ─── Gate Factory ─────────────────────────────────────────────────────────────
+
+
+def pii_postflight_gate(config: PiiGateConfig | None = None) -> "PostflightGate":
+    """Create a PII postflight gate with the given configuration.
+
+    Returns a callable matching the PostflightGate protocol.
+    Zero external dependencies — uses regex patterns from pii_patterns.yaml.
+    """
+    from traceforge.sdk.gate_types import PostflightAction, PostflightVerdict
+
+    cfg = config or PiiGateConfig()
+
+    # Pre-resolve critical entities from YAML if not overridden
+    _ensure_loaded()
+    if cfg.critical_entities is not None:
+        critical_set = frozenset(cfg.critical_entities)
+    else:
+        critical_set = frozenset(e.entity_type for e in _ENTITIES if e.critical)
+
+    def _gate(result: "ToolCallResult", ctx: "GateContext") -> "PostflightVerdict":
+        text = _extract_text(result)
+        if not text or not text.strip():
+            return PostflightVerdict(action=PostflightAction.ACCEPT)
+
+        try:
+            findings = scan_text(
+                text,
+                entities=cfg.entities,
+                score_threshold=cfg.score_threshold,
+                allow_list=frozenset(cfg.allow_list),
+            )
+        except Exception as e:
+            # Fail-closed
+            return PostflightVerdict(
+                action=PostflightAction.SUPPRESS,
+                reason=f"PII gate: scan failed ({type(e).__name__}: {e}) — suppressing",
+            )
+
+        if not findings:
+            return PostflightVerdict(action=PostflightAction.ACCEPT)
+
+        # Check for critical entities
+        if cfg.suppress_on_critical:
+            critical_found = [f for f in findings if f.entity_type in critical_set]
+            if critical_found:
+                types = sorted({f.entity_type for f in critical_found})
+                return PostflightVerdict(
+                    action=PostflightAction.SUPPRESS,
+                    reason=f"Critical PII detected: {', '.join(types)}",
+                )
+
+        # REDACT: use matched text spans as redaction_keys
+        seen: set[str] = set()
+        redaction_keys: list[str] = []
+        for f in sorted(findings, key=lambda x: len(x.text), reverse=True):
+            if f.text not in seen:
+                seen.add(f.text)
+                redaction_keys.append(f.text)
+
+        types_summary = sorted({f.entity_type for f in findings})
+        return PostflightVerdict(
+            action=PostflightAction.REDACT,
+            reason=f"PII detected: {', '.join(types_summary)}",
+            redaction_keys=tuple(redaction_keys),
+        )
+
+    _gate.__qualname__ = "pii_postflight_gate.<locals>._gate"
+    _gate.__doc__ = f"PII postflight gate (threshold={cfg.score_threshold})"
+    return _gate  # type: ignore[return-value]
+
+
+# ─── Helpers ──────────────────────────────────────────────────────────────────
+
+
+def _extract_text(result: "ToolCallResult") -> str:
+    """Extract scannable text from a ToolCallResult."""
+    parts: list[str] = []
+    if result.output:
+        for v in result.output.values():
+            if isinstance(v, str):
+                parts.append(v)
+            elif isinstance(v, (list, tuple)):
+                for item in v:
+                    if isinstance(item, str):
+                        parts.append(item)
+    if result.error:
+        parts.append(result.error)
+    return "\n".join(parts)