traceforge-toolkit 0.1.0__py3-none-any.whl

traceforge/sinks/base.py

@@ -0,0 +1,95 @@
+"""Base storage sink interface."""
+
+from __future__ import annotations
+
+import logging
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, ClassVar
+
+from traceforge.types import SessionEvent, TelemetrySpan, TitleUpdate, UsageRecord
+
+if TYPE_CHECKING:
+    from traceforge.governance.envelope import EnrichedEvent
+
+logger = logging.getLogger(__name__)
+
+
+class StorageSink(ABC):
+    # Sink classes that have already been warned about dropping title updates,
+    # so the warning fires once per class per process rather than per title.
+    _title_drop_warned: ClassVar[set[str]] = set()
+    # Sink classes already warned about dropping non-event envelope payloads
+    # (gap markers / governance-only records) they do not know how to persist.
+    _enriched_drop_warned: ClassVar[set[str]] = set()
+
+    @abstractmethod
+    async def on_event(self, event: SessionEvent) -> None:
+        """Handle a session event. Required for all sinks."""
+        ...
+
+    async def on_enriched_event(self, enriched: "EnrichedEvent") -> None:
+        """Handle a governance-enriched event envelope.
+
+        This is the emission entry point used when governance is wired into the
+        pipeline: the envelope pairs an event with its ``SessionMeta``. The base
+        default keeps every existing sink working unchanged — a live
+        ``SessionEvent`` (which already carries its ``metadata.governance`` stamp)
+        is forwarded to :meth:`on_event`, so a sink that only implements
+        ``on_event`` sees byte-identical output whether or not governance is on.
+
+        Synthetic envelope payloads that are *not* live session events — a
+        :class:`~traceforge.governance.envelope.ContextGapEvent` emitted under
+        backpressure, or a governance-only record — cannot be expressed through
+        ``on_event``. Rather than silently drop them (they signal audit gaps),
+        the base logs a one-time warning naming the sink class. Sinks that want
+        to persist gaps override this method (``JsonlSink``/``SqliteOutputSink``
+        do).
+        """
+        event = enriched.event
+        if isinstance(event, SessionEvent):
+            await self.on_event(event)
+            return
+        cls = type(self).__name__
+        if cls not in StorageSink._enriched_drop_warned:
+            StorageSink._enriched_drop_warned.add(cls)
+            logger.warning(
+                "%s does not handle non-event envelope payloads (e.g. context-gap "
+                "markers); they will be dropped. Override on_enriched_event(enriched) "
+                "to persist or forward them.",
+                cls,
+            )
+
+    async def on_span(self, span: TelemetrySpan) -> None:
+        """Handle a telemetry span. Default no-op."""
+
+    async def on_usage(self, usage: UsageRecord) -> None:
+        """Handle a usage record. Default no-op."""
+
+    async def on_title_update(self, update: TitleUpdate) -> None:
+        """Handle an out-of-band title for a closed activity/step segment.
+
+        Emitted after the segment's events (which already carry its
+        ``activity_id``/``step_id``); keyed to them by ``segment_id``. See
+        :class:`traceforge.types.TitleUpdate`.
+
+        Titles are the structurer's primary output, so — unlike the
+        ``on_span``/``on_usage`` no-ops — the base default does **not** silently
+        drop them: it logs a one-time warning naming the sink class. A custom
+        sink keeps working without overriding this, but the author is told that
+        titles are being discarded and should override to persist them. All
+        in-repo sinks override it.
+        """
+        cls = type(self).__name__
+        if cls not in StorageSink._title_drop_warned:
+            StorageSink._title_drop_warned.add(cls)
+            logger.warning(
+                "%s does not handle title updates; activity/step titles will be "
+                "dropped. Override on_title_update(update) to persist or forward them.",
+                cls,
+            )
+
+    async def flush(self) -> None:
+        """Flush buffered writes. Default no-op."""
+
+    async def close(self) -> None:
+        """Clean up resources. Default no-op."""

traceforge/sinks/callback.py

@@ -0,0 +1,132 @@
+"""Callback-based storage sink for custom event handling."""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+from collections.abc import Awaitable, Callable, Iterable
+
+from traceforge.sinks.base import StorageSink
+from traceforge.types import SessionEvent, TelemetrySpan, TitleUpdate, UsageRecord
+
+#: Accepted shapes for a lightweight event subscriber: a coroutine function or a
+#: plain sync callable (both taking a single :class:`SessionEvent`).
+EventCallback = Callable[[SessionEvent], Awaitable[None] | None]
+
+#: Accepted shapes for a per-subscriber kind filter. ``None`` means "all events".
+#: A string is an exact kind, or a ``"prefix.*"`` wildcard matching that dotted
+#: namespace (e.g. ``"tool.*"`` matches ``tool.call``/``tool.result`` but not
+#: ``tool`` or ``toolbar``). An iterable of such strings matches if *any* does. A
+#: callable is used verbatim as a predicate over the event.
+KindFilter = str | Iterable[str] | Callable[[SessionEvent], bool] | None
+
+
+class CallbackSink(StorageSink):
+    """A sink that delegates to user-provided async callback functions."""
+
+    def __init__(
+        self,
+        on_event: Callable[[SessionEvent], Awaitable[None]] | None = None,
+        on_span: Callable[[TelemetrySpan], Awaitable[None]] | None = None,
+        on_usage: Callable[[UsageRecord], Awaitable[None]] | None = None,
+        on_title_update: Callable[[TitleUpdate], Awaitable[None]] | None = None,
+    ) -> None:
+        self._on_event = on_event
+        self._on_span = on_span
+        self._on_usage = on_usage
+        self._on_title_update = on_title_update
+
+    async def on_event(self, event: SessionEvent) -> None:
+        if self._on_event is not None:
+            await self._on_event(event)
+
+    async def on_span(self, span: TelemetrySpan) -> None:
+        if self._on_span is not None:
+            await self._on_span(span)
+
+    async def on_usage(self, usage: UsageRecord) -> None:
+        if self._on_usage is not None:
+            await self._on_usage(usage)
+
+    async def on_title_update(self, update: TitleUpdate) -> None:
+        if self._on_title_update is not None:
+            await self._on_title_update(update)
+
+
+def _kind_predicate(kind: KindFilter) -> Callable[[SessionEvent], bool] | None:
+    """Compile a :data:`KindFilter` into a predicate over :class:`SessionEvent`.
+
+    Returns ``None`` (meaning "no filtering") for ``None`` or an empty iterable,
+    so the dispatch wrapper can skip the check entirely in the common case.
+    """
+    if kind is None:
+        return None
+    if isinstance(kind, str):
+        patterns: tuple[str, ...] = (kind,)
+    elif callable(kind):
+        return kind
+    else:
+        patterns = tuple(kind)
+        if not patterns:
+            return None
+
+    def predicate(event: SessionEvent) -> bool:
+        return any(_pattern_matches(pattern, event.kind) for pattern in patterns)
+
+    return predicate
+
+
+def _pattern_matches(pattern: str, kind: str) -> bool:
+    """Match one kind string against an exact or ``"prefix.*"`` wildcard pattern."""
+    if pattern.endswith(".*"):
+        # ``"tool.*"`` -> prefix ``"tool."`` -> matches ``tool.call`` but not
+        # bare ``tool`` or an unrelated ``toolbar``.
+        return kind.startswith(pattern[:-1])
+    return kind == pattern
+
+
+def as_async_event_callback(
+    callback: EventCallback,
+    *,
+    kind: KindFilter = None,
+    to_thread: bool = False,
+) -> Callable[[SessionEvent], Awaitable[None]]:
+    """Adapt a sync-or-async event callback into the ``async on_event`` shape.
+
+    This is the adapter behind :meth:`~traceforge.pipeline.EventPipeline.subscribe`:
+    it lets a plain ``Callable[[SessionEvent], None]`` act as a first-class event
+    subscriber without the caller writing ``async def`` or a full sink, and it
+    applies an optional per-subscriber ``kind`` filter *before* dispatch so a
+    filtered-out event never touches the callback.
+
+    - **async callables** (coroutine functions) are awaited directly.
+    - **sync callables** are called inline on the event loop by default — right
+      for the lightweight consumers this exists for (append to a list, put on a
+      queue). Pass ``to_thread=True`` to instead run a blocking sync callback via
+      :func:`asyncio.to_thread`, so it never stalls the loop.
+    - A sync callable that unexpectedly returns an awaitable (e.g. an object whose
+      ``__call__`` is ``async`` and so is not flagged by
+      :func:`asyncio.iscoroutinefunction`) is awaited defensively.
+
+    Raises ``TypeError`` if ``callback`` is not callable.
+    """
+    if not callable(callback):
+        raise TypeError(f"event callback must be callable, got {type(callback).__name__}")
+
+    predicate = _kind_predicate(kind)
+    is_async = asyncio.iscoroutinefunction(callback)
+
+    async def on_event(event: SessionEvent) -> None:
+        if predicate is not None and not predicate(event):
+            return
+        if is_async:
+            await callback(event)
+            return
+        if to_thread:
+            await asyncio.to_thread(callback, event)
+            return
+        result = callback(event)
+        if inspect.isawaitable(result):
+            await result
+
+    return on_event

traceforge/sinks/console.py

@@ -0,0 +1,112 @@
+"""Console sink — pretty-prints governance results to terminal."""
+
+from __future__ import annotations
+
+import logging
+import sys
+
+from traceforge.sinks.base import StorageSink
+from traceforge.types import SessionEvent, TelemetrySpan, TitleUpdate, UsageRecord
+
+logger = logging.getLogger(__name__)
+
+# ANSI color codes
+_RESET = "\033[0m"
+_RED = "\033[91m"
+_YELLOW = "\033[93m"
+_GREEN = "\033[92m"
+_DIM = "\033[2m"
+_BOLD = "\033[1m"
+
+_ACTION_COLORS = {
+    "deny": _RED,
+    "escalate": _RED,
+    "warn": _YELLOW,
+    "allow": _GREEN,
+    "monitor": _DIM,
+}
+
+
+class ConsoleSink(StorageSink):
+    """Prints governance-relevant events to the terminal.
+
+    Only emits events whose governance action matches the configured filter.
+    Designed for real-time human feedback during agent sessions.
+    """
+
+    def __init__(
+        self,
+        filter_actions: list[str] | None = None,
+        color: bool = True,
+        stream: object | None = None,
+    ) -> None:
+        self._filter = set(filter_actions or ["warn", "deny", "escalate"])
+        self._color = color and hasattr(sys.stderr, "isatty") and sys.stderr.isatty()
+        self._stream = stream or sys.stderr
+
+    async def on_event(self, event: SessionEvent) -> None:
+        meta = event.metadata
+        if meta is None:
+            return
+
+        classification = meta.classification
+        if classification is None:
+            return
+
+        # Extract governance action from metadata
+        action = self._extract_action(event)
+        if action is None or action not in self._filter:
+            return
+
+        self._print_event(event, action)
+
+    def _extract_action(self, event: SessionEvent) -> str | None:
+        """Extract the recommended action from event metadata."""
+        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 _print_event(self, event: SessionEvent, action: str) -> None:
+        color = _ACTION_COLORS.get(action, _DIM) if self._color else ""
+        reset = _RESET if self._color else ""
+        bold = _BOLD if self._color else ""
+        dim = _DIM if self._color else ""
+
+        tool_name = event.payload.get("tool_name", event.kind) if event.payload else event.kind
+        args_preview = ""
+        if event.payload:
+            args = event.payload.get("arguments") or event.payload.get("command")
+            if args:
+                args_str = str(args)
+                args_preview = f" {dim}{args_str[:80]}{'...' if len(args_str) > 80 else ''}{reset}"
+
+        risk_score = ""
+        if event.metadata and event.metadata.governance:
+            gov = event.metadata.governance
+            if gov.risk_assessment is not None:
+                risk_score = f" [risk:{gov.risk_assessment.score}]"
+
+        ts = event.timestamp.strftime("%H:%M:%S") if event.timestamp else ""
+        line = f"{dim}{ts}{reset} {color}{bold}{action.upper()}{reset} {tool_name}{risk_score}{args_preview}"
+
+        print(line, file=self._stream)
+
+    async def on_span(self, span: TelemetrySpan) -> None:
+        pass
+
+    async def on_usage(self, usage: UsageRecord) -> None:
+        pass
+
+    async def on_title_update(self, update: TitleUpdate) -> None:
+        """Render a closed segment's title as an indented TOC line."""
+        color = (_GREEN if update.kind == "activity" else _DIM) if self._color else ""
+        reset = _RESET if self._color else ""
+        bold = _BOLD if self._color else ""
+        indent = "" if update.kind == "activity" else "  "
+        label = update.kind.upper()
+        line = f"{indent}{color}{bold}{label}{reset} {update.title}"
+        print(line, file=self._stream)

traceforge/sinks/factory.py

@@ -0,0 +1,94 @@
+"""Single factory mapping declarative ``SinkConfig`` unions to sink instances.
+
+One place builds standard sinks from config so the CLI daemon
+(:func:`traceforge.cli.watch._build_sinks`) and the SDK
+(:meth:`traceforge.sdk.Pipeline.from_config`) stay in parity — the exact
+asymmetry issue #116 closes.
+
+Optional-dependency sinks import lazily *per branch* (e.g. ``boto3`` for S3),
+so a config that never mentions S3 never imports boto3.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from traceforge.config.models import SinkConfig
+    from traceforge.sinks.base import StorageSink
+
+logger = logging.getLogger(__name__)
+
+
+def build_sinks(sink_configs: "list[SinkConfig]") -> "list[StorageSink]":
+    """Instantiate the standard sinks declared by a ``SinkConfig`` list.
+
+    Covers the full serializable union (``sqlite``, ``jsonl``, ``console``,
+    ``webhook``, ``otel``, ``s3``); each config field maps onto its sink
+    constructor faithfully. Callback / custom :class:`StorageSink` subclasses are
+    code-only and intentionally have no ``SinkConfig`` variant, so they are not
+    built here.
+
+    Unknown ``type`` discriminators are logged and skipped. This cannot happen
+    for a validated :data:`~traceforge.config.models.SinkConfig` (the union is
+    exhaustive), but keeping the branch total means a malformed duck-typed config
+    degrades gracefully instead of raising.
+    """
+    sinks: list = []
+    for sink_config in sink_configs:
+        sink_type = getattr(sink_config, "type", None)
+
+        if sink_type == "console":
+            from traceforge.sinks.console import ConsoleSink
+
+            sinks.append(ConsoleSink(filter_actions=sink_config.filter, color=sink_config.color))
+        elif sink_type == "sqlite":
+            from traceforge.sinks.sqlite_output import SqliteOutputSink
+
+            sinks.append(
+                SqliteOutputSink(path=sink_config.path, journal_mode=sink_config.journal_mode)
+            )
+        elif sink_type == "jsonl":
+            from traceforge.sinks.jsonl import JsonlSink
+
+            sinks.append(
+                JsonlSink(path=sink_config.path, rotate_size_mb=sink_config.rotate_size_mb)
+            )
+        elif sink_type == "webhook":
+            from traceforge.sinks.webhook import WebhookSink
+
+            sinks.append(
+                WebhookSink(
+                    url=sink_config.url,
+                    filter_actions=sink_config.filter,
+                    timeout=sink_config.timeout,
+                    max_retries=sink_config.max_retries,
+                    headers=sink_config.headers,
+                )
+            )
+        elif sink_type == "otel":
+            from traceforge.sinks.otel_exporter import OtelExporterSink
+
+            sinks.append(
+                OtelExporterSink(
+                    endpoint=sink_config.endpoint,
+                    service_name=sink_config.service_name,
+                    headers=sink_config.headers,
+                )
+            )
+        elif sink_type == "s3":
+            from traceforge.sinks.s3 import S3Sink
+
+            sinks.append(
+                S3Sink(
+                    bucket=sink_config.bucket,
+                    prefix=sink_config.prefix,
+                    region=sink_config.region,
+                    endpoint_url=sink_config.endpoint_url,
+                )
+            )
+        else:
+            logger.warning("Unknown sink type %r; skipping", sink_type)
+
+    return sinks

traceforge/sinks/jsonl.py

@@ -0,0 +1,125 @@
+"""JSONL file sink — append one JSON line per event."""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+from traceforge.sinks.base import StorageSink
+from traceforge.types import SessionEvent, TelemetrySpan, TitleUpdate, UsageRecord
+
+logger = logging.getLogger(__name__)
+
+
+class JsonlSink(StorageSink):
+    """Appends events as JSON lines to a file.
+
+    Supports ``{session_id}`` template in path for per-session output files.
+    Creates parent directories on first write.
+    """
+
+    def __init__(self, path: str | Path, rotate_size_mb: float | None = None) -> None:
+        self._path_template = str(path)
+        self._rotate_size_mb = rotate_size_mb
+        self._handles: dict[str, object] = {}
+
+    _SAFE_SESSION_RE = __import__("re").compile(r"[^a-zA-Z0-9_\-.]")
+
+    def _resolve_path(self, session_id: str) -> Path:
+        sanitized = self._SAFE_SESSION_RE.sub("_", session_id)[:128]
+        resolved = self._path_template.replace("{session_id}", sanitized)
+        path = Path(resolved).expanduser().resolve()
+        # Ensure resolved path stays under the template's parent directory
+        base = Path(self._path_template.split("{session_id}")[0]).expanduser().resolve()
+        if not str(path).startswith(str(base)):
+            raise ValueError(f"JsonlSink: resolved path escapes base directory: {path}")
+        return path
+
+    async def on_event(self, event: SessionEvent) -> None:
+        path = self._resolve_path(event.session_id)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        line = json.dumps(
+            {
+                "id": event.id,
+                "kind": event.kind,
+                "session_id": event.session_id,
+                "timestamp": event.timestamp.isoformat(),
+                "payload": event.payload,
+                "metadata": event.metadata.model_dump(exclude_none=True)
+                if event.metadata
+                else None,
+            },
+            default=str,
+        )
+
+        try:
+            with open(path, "a", encoding="utf-8") as f:
+                f.write(line + "\n")
+        except OSError as exc:
+            logger.error("JsonlSink: failed to write to %s: %s", path, 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;
+        context-gap markers are written as a tagged ``context_gap`` record."""
+        from traceforge.governance.envelope import ContextGapEvent
+
+        gap = enriched.event
+        if not isinstance(gap, ContextGapEvent):
+            await super().on_enriched_event(enriched)
+            return
+
+        path = self._resolve_path(gap.session_id)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        line = json.dumps(
+            {
+                "record": "context_gap",
+                "id": gap.id,
+                "session_id": gap.session_id,
+                "timestamp": gap.timestamp.isoformat(),
+                "dropped_count": gap.dropped_count,
+                "first_dropped_sequence": gap.first_dropped_sequence,
+                "last_dropped_sequence": gap.last_dropped_sequence,
+                "gap_ordinal": gap.gap_ordinal,
+                "reason": gap.reason,
+                "source_event_key": gap.source_event_key,
+            },
+            default=str,
+        )
+
+        try:
+            with open(path, "a", encoding="utf-8") as f:
+                f.write(line + "\n")
+        except OSError as exc:
+            logger.error("JsonlSink: failed to write to %s: %s", path, exc)
+
+    async def on_title_update(self, update: TitleUpdate) -> None:
+        path = self._resolve_path(update.session_id)
+        path.parent.mkdir(parents=True, exist_ok=True)
+
+        line = json.dumps(
+            {
+                "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,
+            },
+            default=str,
+        )
+
+        try:
+            with open(path, "a", encoding="utf-8") as f:
+                f.write(line + "\n")
+        except OSError as exc:
+            logger.error("JsonlSink: failed to write to %s: %s", path, exc)