proofstep 0.1.0__py3-none-any.whl

proofstep/decorators.py

@@ -0,0 +1,170 @@
+"""Decorators — the lowest-friction way to instrument existing code.
+
+Each inspects the wrapped function and returns a matching sync or async wrapper, so
+the same decorator works on both without the user thinking about it.
+
+`tool` is deliberately separate from `span`. It sets `span_type=tool`, records the
+tool name, and captures the call arguments — which is precisely what the trajectory
+engine consumes. Making it an obvious, distinct decorator is what turns policies from
+a configuration exercise into a one-line change.
+"""
+
+from __future__ import annotations
+
+import functools
+import inspect
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any, Protocol, TypeVar
+
+from proofstep.safety import NOOP
+from proofstep_types import SpanType
+
+if TYPE_CHECKING:
+    from proofstep.client import Client
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+
+class Decorator(Protocol):
+    """A decorator factory that preserves the wrapped function's signature.
+
+    Positional-or-keyword `name` plus arbitrary keywords: `span` accepts several
+    extra keyword-only options that callers may pass through.
+    """
+
+    def __call__(self, name: str | None = ..., /, **kwargs: Any) -> Callable[[F], F]: ...
+
+
+class SpanDecorator(Protocol):
+    """`span` exposes its options explicitly rather than behind `**kwargs`.
+
+    A typo in `capture_args` should be a type error, not a silently ignored keyword.
+    """
+
+    def __call__(
+        self,
+        name: str | None = ...,
+        /,
+        *,
+        span_type: SpanType | str = ...,
+        tool_name: str | None = ...,
+        capture_args: bool = ...,
+        capture_result: bool = ...,
+    ) -> Callable[[F], F]: ...
+
+
+# Arguments that are never worth capturing: they are either huge or self-referential.
+_SKIP_ARGS = frozenset({"self", "cls"})
+
+
+def _bind_args(
+    fn: Callable[..., Any], args: tuple[Any, ...], kwargs: dict[str, Any]
+) -> dict[str, Any]:
+    """Best-effort mapping of a call to a plain dict for `tool_args`."""
+    try:
+        bound = inspect.signature(fn).bind_partial(*args, **kwargs)
+        bound.apply_defaults()
+        return {k: v for k, v in bound.arguments.items() if k not in _SKIP_ARGS}
+    except (TypeError, ValueError):
+        return {"args": list(args), "kwargs": kwargs}
+
+
+def make_trace(client_of: Callable[[], Client]) -> Decorator:
+    def trace(name: str | None = None, /, **trace_kwargs: Any) -> Callable[[F], F]:
+        def decorate(fn: F) -> F:
+            label = name or fn.__name__
+
+            if inspect.iscoroutinefunction(fn):
+
+                @functools.wraps(fn)
+                async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                    with client_of().trace(label, **trace_kwargs) as recorder:
+                        result = await fn(*args, **kwargs)
+                        if recorder is not NOOP:
+                            recorder.set_metadata(**trace_kwargs.get("metadata", {}))
+                        return result
+
+                return async_wrapper  # type: ignore[return-value]
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with client_of().trace(label, **trace_kwargs):
+                    return fn(*args, **kwargs)
+
+            return sync_wrapper  # type: ignore[return-value]
+
+        return decorate
+
+    return trace
+
+
+def make_span(client_of: Callable[[], Client]) -> SpanDecorator:
+    def span(
+        name: str | None = None,
+        /,
+        *,
+        span_type: SpanType | str = SpanType.CUSTOM,
+        tool_name: str | None = None,
+        capture_args: bool = False,
+        capture_result: bool = True,
+    ) -> Callable[[F], F]:
+        def decorate(fn: F) -> F:
+            label = name or fn.__name__
+
+            def before(recorder: Any, args: tuple[Any, ...], kwargs: dict[str, Any]) -> None:
+                if recorder is NOOP:
+                    return
+                bound = _bind_args(fn, args, kwargs)
+                if capture_args:
+                    recorder.set_args(bound)
+                recorder.set_input(bound)
+
+            def after(recorder: Any, result: Any) -> None:
+                if recorder is not NOOP and capture_result:
+                    recorder.set_output(result)
+
+            if inspect.iscoroutinefunction(fn):
+
+                @functools.wraps(fn)
+                async def async_wrapper(*args: Any, **kwargs: Any) -> Any:
+                    with client_of().span(label, span_type=span_type, tool_name=tool_name) as rec:
+                        before(rec, args, kwargs)
+                        result = await fn(*args, **kwargs)
+                        after(rec, result)
+                        return result
+
+                return async_wrapper  # type: ignore[return-value]
+
+            @functools.wraps(fn)
+            def sync_wrapper(*args: Any, **kwargs: Any) -> Any:
+                with client_of().span(label, span_type=span_type, tool_name=tool_name) as rec:
+                    before(rec, args, kwargs)
+                    result = fn(*args, **kwargs)
+                    after(rec, result)
+                    return result
+
+            return sync_wrapper  # type: ignore[return-value]
+
+        return decorate
+
+    return span
+
+
+def make_tool(span_factory: SpanDecorator) -> Decorator:
+    def tool(name: str | None = None, /, **kwargs: Any) -> Callable[[F], F]:
+        """Instrument a tool call. Captures arguments, which policies read."""
+
+        def decorate(fn: F) -> F:
+            label = name or fn.__name__
+            decorated: F = span_factory(
+                label,
+                span_type=SpanType.TOOL,
+                tool_name=label,
+                capture_args=True,
+                **kwargs,
+            )(fn)
+            return decorated
+
+        return decorate
+
+    return tool

proofstep/exporter.py

@@ -0,0 +1,276 @@
+"""Batching exporter with a bounded buffer.
+
+The contract with the host application, in priority order:
+
+1. **Never block.** Submitting a trace is a non-blocking enqueue. If the buffer is
+   full we drop and count; a full buffer must never become backpressure on the
+   caller's request path.
+2. **Never raise.** Every failure is swallowed and logged once per window.
+3. **Lose visibly.** Drops increment a counter that is reported on the trace itself.
+   Silent loss is worse than visible loss, because it looks like the workflow simply
+   did not do the thing.
+
+A background thread rather than an asyncio task: the SDK must work identically in a
+sync Django view, a Celery worker, and an async FastAPI handler, and a thread is the
+only one of those that needs no event loop.
+"""
+
+from __future__ import annotations
+
+import gzip
+import json
+import queue
+import random
+import threading
+import time
+from pathlib import Path
+from typing import TYPE_CHECKING, Any
+
+from proofstep.safety import log_once, never_raises
+
+if TYPE_CHECKING:
+    from proofstep.config import Config
+    from proofstep_types import Trace
+
+USER_AGENT = "proofstep-python"
+
+
+class ExportStats:
+    """Observable counters. A telemetry client that cannot report its own health is
+    the one thing worse than no telemetry."""
+
+    __slots__ = ("dropped_traces", "exported_spans", "exported_traces", "failures", "spooled")
+
+    def __init__(self) -> None:
+        self.exported_traces = 0
+        self.exported_spans = 0
+        self.dropped_traces = 0
+        self.failures = 0
+        self.spooled = 0
+
+    def as_dict(self) -> dict[str, int]:
+        return {slot: getattr(self, slot) for slot in self.__slots__}
+
+
+class Exporter:
+    def __init__(self, config: Config, *, transport: Any = None) -> None:
+        self.config = config
+        self.stats = ExportStats()
+        self._queue: queue.Queue[Trace | None] = queue.Queue(maxsize=_capacity(config))
+        self._transport = transport
+        self._thread: threading.Thread | None = None
+        self._stopping = threading.Event()
+        self._idle = threading.Event()
+        self._idle.set()
+        self._lock = threading.Lock()
+
+    # ------------------------------------------------------------------ lifecycle
+
+    def _ensure_started(self) -> None:
+        if self._thread is not None and self._thread.is_alive():
+            return
+        with self._lock:
+            if self._thread is not None and self._thread.is_alive():
+                return
+            self._stopping.clear()
+            # Daemon: a forgotten flush() must never stop the process exiting.
+            # `atexit` handles the orderly case.
+            self._thread = threading.Thread(
+                target=self._run, name="proofstep-exporter", daemon=True
+            )
+            self._thread.start()
+
+    @never_raises()
+    def submit(self, trace: Trace) -> None:
+        """Non-blocking enqueue. Drops the oldest item when full."""
+        if not self.config.sends:
+            return
+        self._ensure_started()
+        self._idle.clear()
+        try:
+            self._queue.put_nowait(trace)
+        except queue.Full:
+            # Drop the oldest: the newest trace is the one someone is probably
+            # watching for, and an unbounded queue is how a telemetry client turns
+            # a provider outage into an OOM.
+            try:
+                self._queue.get_nowait()
+                self.stats.dropped_traces += 1
+                self._queue.put_nowait(trace)
+            except (queue.Empty, queue.Full):
+                self.stats.dropped_traces += 1
+            log_once(
+                "exporter.full",
+                f"export buffer is full ({self._queue.maxsize} traces); dropping oldest. "
+                f"{self.stats.dropped_traces} dropped so far.",
+            )
+
+    @never_raises(default=False)
+    def flush(self, timeout: float | None = None) -> bool:
+        """Wait for the queue to drain. Returns False on timeout."""
+        if self._thread is None or not self._thread.is_alive():
+            return True
+        deadline = time.monotonic() + (
+            timeout if timeout is not None else self.config.shutdown_timeout_s
+        )
+        while time.monotonic() < deadline:
+            if self._queue.empty() and self._idle.is_set():
+                return True
+            time.sleep(0.01)
+        log_once("exporter.flush_timeout", "flush timed out; some traces were not sent")
+        return False
+
+    @never_raises()
+    def shutdown(self, timeout: float | None = None) -> None:
+        self.flush(timeout)
+        self._stopping.set()
+        with self._lock:
+            thread, self._thread = self._thread, None
+        if thread is not None and thread.is_alive():
+            self._queue.put_nowait(None)  # wake the worker
+            thread.join(timeout=1.0)
+
+    # --------------------------------------------------------------------- worker
+
+    def _run(self) -> None:
+        batch: list[Trace] = []
+        last_flush = time.monotonic()
+
+        while not self._stopping.is_set():
+            timeout = max(0.01, self.config.flush_interval_s - (time.monotonic() - last_flush))
+            try:
+                item = self._queue.get(timeout=timeout)
+            except queue.Empty:
+                item = None
+            else:
+                if item is None:
+                    break
+                batch.append(item)
+
+            due = (time.monotonic() - last_flush) >= self.config.flush_interval_s
+            spans = sum(len(t.spans) for t in batch)
+            if batch and (due or spans >= self.config.batch_size):
+                self._send_batch(batch)
+                batch = []
+                last_flush = time.monotonic()
+
+            if not batch and self._queue.empty():
+                self._idle.set()
+
+        if batch:
+            self._send_batch(batch)
+        self._idle.set()
+
+    def _send_batch(self, batch: list[Trace]) -> None:
+        payload = _encode(batch, self.config)
+        if payload is None:
+            return
+
+        body = gzip.compress(payload)
+        for attempt in range(self.config.max_retries):
+            if self._stopping.is_set() and attempt > 0:
+                break
+            try:
+                self._post(body)
+            except Exception as exc:
+                self.stats.failures += 1
+                if attempt == self.config.max_retries - 1:
+                    log_once(
+                        "exporter.send_failed",
+                        f"could not reach {self.config.endpoint} after "
+                        f"{self.config.max_retries} attempts ({type(exc).__name__}: {exc}); "
+                        "traces are being spooled or dropped",
+                    )
+                    self._spool(body)
+                    return
+                # Exponential backoff with full jitter: synchronized retries from
+                # many processes are how a struggling server stays down.
+                delay = min(30.0, (2**attempt) * 0.5)
+                time.sleep(random.uniform(0, delay))  # noqa: S311 — jitter, not crypto
+            else:
+                self.stats.exported_traces += len(batch)
+                self.stats.exported_spans += sum(len(t.spans) for t in batch)
+                return
+
+    def _post(self, body: bytes) -> None:
+        if self._transport is not None:
+            self._transport(body)
+            return
+
+        import httpx  # noqa: PLC0415 — lazy so importing the SDK costs nothing
+
+        response = httpx.post(
+            f"{self.config.endpoint}/v1/ingest/traces",
+            content=body,
+            headers={
+                "Authorization": f"Bearer {self.config.api_key}",
+                "Content-Type": "application/json",
+                "Content-Encoding": "gzip",
+                "User-Agent": USER_AGENT,
+            },
+            timeout=self.config.export_timeout_s,
+        )
+        if response.status_code >= 500 or response.status_code == 429:
+            msg = f"server returned {response.status_code}"
+            raise RuntimeError(msg)
+        if response.status_code >= 400:
+            # 4xx is our bug or a bad key. Retrying cannot fix it, and retrying an
+            # auth failure five times per batch is how you get rate-limited too.
+            log_once(
+                "exporter.rejected",
+                f"ingest rejected the batch with {response.status_code}: {response.text[:200]}",
+            )
+
+    def _spool(self, body: bytes) -> None:
+        if self.config.spool_dir is None:
+            return
+        try:
+            directory = Path(self.config.spool_dir)
+            directory.mkdir(parents=True, exist_ok=True)
+            name = f"{int(time.time() * 1000)}-{random.randbytes(4).hex()}.json.gz"  # noqa: S311
+            (directory / name).write_bytes(body)
+            self.stats.spooled += 1
+        except OSError as exc:
+            log_once("exporter.spool_failed", f"could not spool batch: {exc}")
+
+
+def _capacity(config: Config) -> int:
+    """Queue length in traces, derived from the span budget."""
+    return max(16, config.max_buffered_spans // 20)
+
+
+def _encode(batch: list[Trace], config: Config) -> bytes | None:
+    traces: list[dict[str, Any]] = []
+    spans: list[dict[str, Any]] = []
+    for trace in batch:
+        dumped = trace.model_dump(mode="json", exclude={"spans"})
+        traces.append(dumped)
+        for span in trace.spans:
+            spans.append(span.model_dump(mode="json"))
+
+    payload = {
+        "resource": {
+            "service.name": config.service_name or config.project or "unknown",
+            "environment": config.environment,
+            "git.commit": config.git_commit,
+            "sdk.name": USER_AGENT,
+        },
+        "traces": traces,
+        "spans": spans,
+        "dropped_span_count": sum(t.dropped_span_count for t in batch),
+    }
+
+    try:
+        encoded = json.dumps(payload, default=str).encode("utf-8")
+    except (TypeError, ValueError) as exc:
+        log_once("exporter.encode_failed", f"could not serialize batch: {exc}")
+        return None
+
+    if len(encoded) > config.max_batch_bytes:
+        log_once(
+            "exporter.batch_too_large",
+            f"batch of {len(encoded)} bytes exceeds max_batch_bytes "
+            f"({config.max_batch_bytes}); dropping it",
+        )
+        return None
+    return encoded

proofstep/propagation.py

@@ -0,0 +1,49 @@
+"""W3C trace-context propagation across service boundaries.
+
+Uses the standard `traceparent` header rather than a bespoke one, so a request that
+crosses into a service instrumented with plain OpenTelemetry still stitches together.
+"""
+
+from __future__ import annotations
+
+import re
+from typing import TYPE_CHECKING
+
+from proofstep.context import current_span, current_trace
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, MutableMapping
+
+TRACEPARENT = "traceparent"
+_FORMAT = re.compile(r"^00-([0-9a-f]{32})-([0-9a-f]{16})-([0-9a-f]{2})$")
+
+
+def inject(headers: MutableMapping[str, str] | None = None) -> dict[str, str]:
+    """Add `traceparent` for the current span, if any."""
+    out: dict[str, str] = dict(headers or {})
+    trace = current_trace()
+    span = current_span()
+    if trace is None or span is None:
+        return out
+    flags = "01" if trace.sampled else "00"
+    out[TRACEPARENT] = f"00-{trace.trace_id}-{span.span_id}-{flags}"
+    return out
+
+
+def extract(headers: Mapping[str, str]) -> tuple[str, str, bool] | None:
+    """Parse `traceparent` into (trace_id, parent_span_id, sampled).
+
+    Returns None on anything malformed. A caller-supplied header is untrusted input
+    and must never be able to corrupt our ids.
+    """
+    lowered = {k.lower(): v for k, v in headers.items()}
+    raw = lowered.get(TRACEPARENT)
+    if not raw:
+        return None
+    match = _FORMAT.match(raw.strip())
+    if not match:
+        return None
+    trace_id, span_id, flags = match.groups()
+    if trace_id == "0" * 32 or span_id == "0" * 16:
+        return None
+    return trace_id, span_id, bool(int(flags, 16) & 0x01)