traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/sinks/sqlite_output.py

@@ -0,0 +1,234 @@
+"""SQLite output sink — stores enriched events in a queryable database."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+import sqlite3
+from pathlib import Path
+
+from traceforge.sinks.base import StorageSink
+from traceforge.types import SessionEvent, TelemetrySpan, TitleUpdate, UsageRecord
+
+logger = logging.getLogger(__name__)
+
+_CREATE_TABLE = """\
+CREATE TABLE IF NOT EXISTS enriched_events (
+    id TEXT PRIMARY KEY,
+    session_id TEXT NOT NULL,
+    kind TEXT NOT NULL,
+    timestamp TEXT NOT NULL,
+    tool_name TEXT,
+    risk_level TEXT,
+    risk_score INTEGER,
+    action TEXT,
+    payload_json TEXT,
+    metadata_json TEXT,
+    created_at TEXT NOT NULL DEFAULT (datetime('now'))
+);
+
+CREATE INDEX IF NOT EXISTS idx_enriched_session ON enriched_events(session_id);
+CREATE INDEX IF NOT EXISTS idx_enriched_timestamp ON enriched_events(timestamp);
+CREATE INDEX IF NOT EXISTS idx_enriched_risk ON enriched_events(risk_level);
+CREATE INDEX IF NOT EXISTS idx_enriched_action ON enriched_events(action);
+
+CREATE TABLE IF NOT EXISTS segment_titles (
+    segment_id TEXT NOT NULL,
+    kind TEXT NOT NULL,
+    session_id TEXT NOT NULL,
+    title TEXT NOT NULL,
+    version INTEGER NOT NULL,
+    parent_id TEXT,
+    updated_at TEXT NOT NULL DEFAULT (datetime('now')),
+    PRIMARY KEY (segment_id, kind)
+);
+
+CREATE INDEX IF NOT EXISTS idx_segment_titles_session ON segment_titles(session_id);
+
+CREATE TABLE IF NOT EXISTS context_gaps (
+    id TEXT PRIMARY KEY,
+    session_id TEXT NOT NULL,
+    timestamp TEXT NOT NULL,
+    dropped_count INTEGER NOT NULL,
+    first_dropped_sequence INTEGER,
+    last_dropped_sequence INTEGER,
+    gap_ordinal INTEGER NOT NULL,
+    reason TEXT NOT NULL,
+    source_event_key TEXT NOT NULL,
+    created_at TEXT NOT NULL DEFAULT (datetime('now'))
+);
+
+CREATE INDEX IF NOT EXISTS idx_context_gaps_session ON context_gaps(session_id);
+"""
+
+
+class SqliteOutputSink(StorageSink):
+    """Persists enriched events to a SQLite database.
+
+    Provides a queryable local audit trail with indexed columns for
+    session_id, timestamp, risk_level, and action.
+    """
+
+    _VALID_JOURNAL_MODES = frozenset({"wal", "delete", "truncate", "persist", "memory", "off"})
+
+    def __init__(self, path: str | Path, journal_mode: str = "wal") -> None:
+        self._path = Path(path).expanduser()
+        if journal_mode.lower() not in self._VALID_JOURNAL_MODES:
+            raise ValueError(
+                f"Invalid journal_mode '{journal_mode}'. "
+                f"Must be one of: {', '.join(sorted(self._VALID_JOURNAL_MODES))}"
+            )
+        self._journal_mode = journal_mode.lower()
+        self._conn: sqlite3.Connection | None = None
+
+    async def on_event(self, event: SessionEvent) -> None:
+        tool_name = None
+        risk_level = None
+        risk_score = None
+        action = None
+
+        if event.payload:
+            tool_name = event.payload.get("tool_name")
+
+        if event.metadata:
+            gov = event.metadata.governance
+            if gov is not None:
+                if gov.risk_assessment is not None:
+                    risk_level = gov.risk_assessment.level
+                    risk_score = gov.risk_assessment.score
+                if gov.recommendation is not None:
+                    action = gov.recommendation.recommended_action.value
+
+        payload_json = json.dumps(event.payload, default=str) if event.payload else None
+        metadata_json = (
+            json.dumps(event.metadata.model_dump(exclude_none=True), default=str)
+            if event.metadata
+            else None
+        )
+
+        params = (
+            event.id,
+            event.session_id,
+            event.kind,
+            event.timestamp.isoformat() if event.timestamp else None,
+            tool_name,
+            risk_level,
+            risk_score,
+            action,
+            payload_json,
+            metadata_json,
+        )
+
+        await asyncio.to_thread(self._write_event, params)
+
+    def _write_event(self, params: tuple) -> None:
+        """Synchronous write — called via asyncio.to_thread."""
+        conn = self._get_conn()
+        try:
+            conn.execute(
+                """INSERT OR IGNORE INTO enriched_events
+                   (id, session_id, kind, timestamp, tool_name, risk_level, risk_score, action, payload_json, metadata_json)
+                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+                params,
+            )
+            conn.commit()
+        except sqlite3.Error as exc:
+            logger.error("SqliteOutputSink: write failed: %s", exc)
+
+    async def on_span(self, span: TelemetrySpan) -> None:
+        pass
+
+    async def on_usage(self, usage: UsageRecord) -> None:
+        pass
+
+    async def on_enriched_event(self, enriched) -> None:
+        """Persist a governance envelope. Live events keep byte-identical output
+        (delegated to :meth:`on_event`); context-gap markers land in the
+        ``context_gaps`` table."""
+        from traceforge.governance.envelope import ContextGapEvent
+
+        gap = enriched.event
+        if not isinstance(gap, ContextGapEvent):
+            await super().on_enriched_event(enriched)
+            return
+
+        params = (
+            gap.id,
+            gap.session_id,
+            gap.timestamp.isoformat() if gap.timestamp else None,
+            gap.dropped_count,
+            gap.first_dropped_sequence,
+            gap.last_dropped_sequence,
+            gap.gap_ordinal,
+            gap.reason,
+            gap.source_event_key,
+        )
+        await asyncio.to_thread(self._write_gap, params)
+
+    def _write_gap(self, params: tuple) -> None:
+        """Synchronous gap write — called via asyncio.to_thread."""
+        conn = self._get_conn()
+        try:
+            conn.execute(
+                """INSERT OR IGNORE INTO context_gaps
+                   (id, session_id, timestamp, dropped_count, first_dropped_sequence,
+                    last_dropped_sequence, gap_ordinal, reason, source_event_key)
+                   VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)""",
+                params,
+            )
+            conn.commit()
+        except sqlite3.Error as exc:
+            logger.error("SqliteOutputSink: gap write failed: %s", exc)
+
+    async def on_title_update(self, update: TitleUpdate) -> None:
+        params = (
+            update.segment_id,
+            update.kind,
+            update.session_id,
+            update.title,
+            update.version,
+            update.parent_id,
+        )
+        await asyncio.to_thread(self._write_title, params)
+
+    def _write_title(self, params: tuple) -> None:
+        """Synchronous upsert keeping the highest version — called via to_thread."""
+        conn = self._get_conn()
+        try:
+            conn.execute(
+                """INSERT INTO segment_titles
+                   (segment_id, kind, session_id, title, version, parent_id)
+                   VALUES (?, ?, ?, ?, ?, ?)
+                   ON CONFLICT(segment_id, kind) DO UPDATE SET
+                       title=excluded.title,
+                       session_id=excluded.session_id,
+                       version=excluded.version,
+                       parent_id=excluded.parent_id,
+                       updated_at=datetime('now')
+                   WHERE excluded.version >= segment_titles.version""",
+                params,
+            )
+            conn.commit()
+        except sqlite3.Error as exc:
+            logger.error("SqliteOutputSink: title write failed: %s", exc)
+
+    async def flush(self) -> None:
+        if self._conn:
+            self._conn.commit()
+
+    async def close(self) -> None:
+        if self._conn:
+            self._conn.close()
+            self._conn = None
+
+    def _get_conn(self) -> sqlite3.Connection:
+        if self._conn is None:
+            self._path.parent.mkdir(parents=True, exist_ok=True)
+            self._conn = sqlite3.connect(str(self._path), timeout=5, check_same_thread=False)
+            try:
+                self._conn.execute(f"PRAGMA journal_mode={self._journal_mode}")
+            except sqlite3.OperationalError:
+                pass
+            self._conn.executescript(_CREATE_TABLE)
+        return self._conn

traceforge/sinks/webhook.py

@@ -0,0 +1,136 @@
+"""Webhook sink — POST governance results to an HTTP endpoint."""
+
+from __future__ import annotations
+
+import asyncio
+import json
+import logging
+from urllib.error import URLError
+from urllib.request import Request, urlopen
+
+from traceforge.sinks.base import StorageSink
+from traceforge.types import SessionEvent, TelemetrySpan, TitleUpdate, UsageRecord
+
+logger = logging.getLogger(__name__)
+
+
+class WebhookSink(StorageSink):
+    """POSTs enriched events as JSON to a configured URL.
+
+    Only emits events matching the filter (by governance action).
+    Uses stdlib urllib to avoid adding dependencies.
+    """
+
+    def __init__(
+        self,
+        url: str,
+        filter_actions: list[str] | None = None,
+        timeout: float = 10.0,
+        max_retries: int = 3,
+        headers: dict[str, str] | None = None,
+    ) -> None:
+        if not url.startswith(("http://", "https://")):
+            raise ValueError(f"Webhook URL must use http:// or https:// scheme, got: {url}")
+        self._url = url
+        self._filter = set(filter_actions or ["deny", "escalate"])
+        self._timeout = timeout
+        self._max_retries = max_retries
+        self._headers = headers or {}
+
+    async def on_event(self, event: SessionEvent) -> None:
+        action = self._extract_action(event)
+        if action is None or action not in self._filter:
+            return
+
+        payload = {
+            "id": event.id,
+            "kind": event.kind,
+            "session_id": event.session_id,
+            "timestamp": event.timestamp.isoformat() if event.timestamp else None,
+            "payload": event.payload,
+            "governance": self._extract_governance(event),
+        }
+        await self._post(payload)
+
+    async def on_title_update(self, update: TitleUpdate) -> None:
+        await self._post(
+            {
+                "record": "title_update",
+                "session_id": update.session_id,
+                "segment_id": update.segment_id,
+                "kind": update.kind,
+                "title": update.title,
+                "version": update.version,
+                "parent_id": update.parent_id,
+            }
+        )
+
+    async def _post(self, payload: dict) -> None:
+        body = json.dumps(payload, default=str).encode("utf-8")
+        headers = {
+            "Content-Type": "application/json",
+            **self._headers,
+        }
+
+        for attempt in range(1, self._max_retries + 1):
+            try:
+                req = Request(self._url, data=body, headers=headers, method="POST")
+                resp = await asyncio.to_thread(self._do_request, req)
+                if resp < 300:
+                    return
+                logger.warning(
+                    "WebhookSink: %s returned status %d (attempt %d/%d)",
+                    self._url,
+                    resp,
+                    attempt,
+                    self._max_retries,
+                )
+            except (URLError, OSError, TimeoutError) as exc:
+                logger.warning(
+                    "WebhookSink: POST to %s failed (attempt %d/%d): %s",
+                    self._url,
+                    attempt,
+                    self._max_retries,
+                    exc,
+                )
+
+        logger.error("WebhookSink: all %d attempts to %s failed", self._max_retries, self._url)
+
+    def _do_request(self, req: Request) -> int:
+        """Synchronous HTTP request — returns status code. Ensures response body is consumed."""
+        with urlopen(req, timeout=self._timeout) as resp:
+            resp.read()  # consume body to release connection
+            return resp.status
+
+    def _extract_action(self, event: SessionEvent) -> str | None:
+        gov = event.metadata.governance if event.metadata else None
+        if gov is None:
+            return None
+        rec = gov.recommendation
+        if rec is None:
+            return None
+        return rec.recommended_action.value
+
+    def _extract_governance(self, event: SessionEvent) -> dict | None:
+        gov = event.metadata.governance if event.metadata else None
+        if gov is None:
+            return None
+        result: dict = {}
+        if gov.risk_assessment is not None:
+            result["risk_assessment"] = {
+                "score": gov.risk_assessment.score,
+                "level": gov.risk_assessment.level,
+                "confidence": gov.risk_assessment.confidence,
+            }
+        if gov.recommendation is not None:
+            result["recommendation"] = {
+                "action": gov.recommendation.recommended_action.value,
+                "reason_code": gov.recommendation.reason_code,
+            }
+        return result or None
+
+    async def on_span(self, span: TelemetrySpan) -> None:
+        pass
+
+    async def on_usage(self, usage: UsageRecord) -> None:
+        pass

traceforge/sources/__init__.py

@@ -0,0 +1,18 @@
+from traceforge.sources.base import RawRecord, Source
+from traceforge.sources.file_poll import FilePollSource
+from traceforge.sources.file_watch import FileWatchSource
+from traceforge.sources.http_poll import HttpPollSource
+from traceforge.sources.replay import ReplaySource
+from traceforge.sources.sqlite import SqliteSource
+from traceforge.sources.sse import SSESource
+
+__all__ = [
+    "FilePollSource",
+    "FileWatchSource",
+    "HttpPollSource",
+    "RawRecord",
+    "ReplaySource",
+    "SqliteSource",
+    "SSESource",
+    "Source",
+]

traceforge/sources/auto_detect.py

@@ -0,0 +1,173 @@
+"""Framework auto-detection — discovers installed AI coding agents by checking well-known paths."""
+
+from __future__ import annotations
+
+import logging
+import os
+import platform
+from pathlib import Path
+
+logger = logging.getLogger(__name__)
+
+
+def _expand(path: str) -> Path:
+    return Path(os.path.expanduser(path))
+
+
+def _vscode_global_storage() -> Path | None:
+    """Locate VS Code's globalStorage directory (OS-specific)."""
+    system = platform.system()
+    if system == "Darwin":
+        base = _expand("~/Library/Application Support/Code/User/globalStorage")
+    elif system == "Windows":
+        appdata = os.environ.get("APPDATA", "")
+        base = Path(appdata) / "Code" / "User" / "globalStorage" if appdata else None
+        if base is None:
+            return None
+    else:  # Linux
+        base = _expand("~/.config/Code/User/globalStorage")
+    return base if base.is_dir() else None
+
+
+def _goose_data_dir() -> Path:
+    """Goose uses XDG data dir conventions via the `etcetera` crate."""
+    system = platform.system()
+    if system == "Darwin":
+        return _expand("~/Library/Application Support/Block/goose")
+    elif system == "Windows":
+        local = os.environ.get("LOCALAPPDATA", "")
+        return Path(local) / "Block" / "goose" if local else _expand("~/.local/share/goose")
+    else:
+        return _expand("~/.local/share/goose")
+
+
+def _amazonq_data_dir() -> Path:
+    """Amazon Q Developer CLI data directory."""
+    system = platform.system()
+    if system == "Darwin":
+        return _expand("~/Library/Application Support/amazon-q")
+    elif system == "Windows":
+        local = os.environ.get("LOCALAPPDATA", "")
+        return Path(local) / "amazon-q" if local else _expand("~/.local/share/amazon-q")
+    else:
+        return _expand("~/.local/share/amazon-q")
+
+
+# ─── Detection registry ─────────────────────────────────────────────────────
+
+
+class DetectedFramework:
+    """A detected framework with its source path and adapter config."""
+
+    def __init__(self, name: str, path: Path, adapter: str, ingestion_mode: str) -> None:
+        self.name = name
+        self.path = path
+        self.adapter = adapter
+        self.ingestion_mode = ingestion_mode
+
+    def __repr__(self) -> str:
+        return f"DetectedFramework({self.name!r}, path={self.path})"
+
+
+def detect_frameworks(frameworks: list[str] | None = None) -> list[DetectedFramework]:
+    """Scan well-known paths for installed AI coding agent frameworks.
+
+    Args:
+        frameworks: If provided, only detect these specific frameworks.
+                   If None/empty, detect all known frameworks.
+
+    Returns:
+        List of detected frameworks with their paths and adapter configs.
+    """
+    all_detectors = {
+        "claude": _detect_claude,
+        "codex": _detect_codex,
+        "continue": _detect_continue,
+        "cline": _detect_cline,
+        "goose": _detect_goose,
+        "amazonq": _detect_amazonq,
+        "aider": _detect_aider,
+    }
+
+    target = frameworks if frameworks else list(all_detectors.keys())
+    detected: list[DetectedFramework] = []
+
+    for name in target:
+        detector = all_detectors.get(name)
+        if detector is None:
+            logger.debug("Unknown framework for detection: %s", name)
+            continue
+        try:
+            result = detector()
+            if result is not None:
+                detected.append(result)
+                logger.info("Detected %s at %s", result.name, result.path)
+        except Exception as exc:
+            logger.debug("Detection failed for %s: %s", name, exc)
+
+    return detected
+
+
+# ─── Individual detectors ────────────────────────────────────────────────────
+
+
+def _detect_claude() -> DetectedFramework | None:
+    path = _expand("~/.claude/projects")
+    if path.is_dir():
+        return DetectedFramework("claude", path, "claude", "file_watch")
+    return None
+
+
+def _detect_codex() -> DetectedFramework | None:
+    codex_home = os.environ.get("CODEX_HOME")
+    if codex_home:
+        path = Path(codex_home) / "sessions"
+    else:
+        path = _expand("~/.codex/sessions")
+    if path.is_dir():
+        return DetectedFramework("codex", path, "codex", "file_watch")
+    return None
+
+
+def _detect_continue() -> DetectedFramework | None:
+    custom = os.environ.get("CONTINUE_GLOBAL_DIR")
+    if custom:
+        path = Path(custom) / "sessions"
+    else:
+        path = _expand("~/.continue/sessions")
+    if path.is_dir():
+        return DetectedFramework("continue", path, "continue", "file_watch")
+    return None
+
+
+def _detect_cline() -> DetectedFramework | None:
+    gs = _vscode_global_storage()
+    if gs is None:
+        return None
+    path = gs / "saoudrizwan.claude-dev" / "tasks"
+    if path.is_dir():
+        return DetectedFramework("cline", path, "cline", "file_watch")
+    return None
+
+
+def _detect_goose() -> DetectedFramework | None:
+    path = _goose_data_dir() / "sessions"
+    if path.is_dir():
+        return DetectedFramework("goose", path, "goose", "poll")
+    return None
+
+
+def _detect_amazonq() -> DetectedFramework | None:
+    path = _amazonq_data_dir() / "data.sqlite3"
+    if path.is_file():
+        return DetectedFramework("amazonq", path, "amazonq", "sqlite")
+    return None
+
+
+def _detect_aider() -> DetectedFramework | None:
+    # Aider writes .aider.chat.history.md in the project root
+    # We check CWD since aider is project-local
+    path = Path.cwd() / ".aider.chat.history.md"
+    if path.is_file():
+        return DetectedFramework("aider", path, "aider_markdown", "file_watch")
+    return None

traceforge/sources/base.py

@@ -0,0 +1,45 @@
+"""Base abstractions for configurable ingestion sources."""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import AsyncIterator
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from types import TracebackType
+
+from traceforge.types import IngestionMode
+
+
+@dataclass(slots=True)
+class RawRecord:
+    """A complete raw record emitted by a source transport."""
+
+    payload: str
+    source_name: str
+    mode: IngestionMode
+    sequence: int | None = None
+    received_at: datetime = field(default_factory=lambda: datetime.now(timezone.utc))
+
+
+class Source(ABC):
+    """Async source that yields complete logical records."""
+
+    name: str
+
+    @abstractmethod
+    async def __aenter__(self) -> "Source":
+        """Allocate any resources required by the source."""
+
+    @abstractmethod
+    async def __aexit__(
+        self,
+        exc_type: type[BaseException] | None,
+        exc: BaseException | None,
+        tb: TracebackType | None,
+    ) -> None:
+        """Release any resources held by the source."""
+
+    @abstractmethod
+    def __aiter__(self) -> AsyncIterator[RawRecord]:
+        """Yield complete logical records from the source."""