agentrail-python 0.1.0__py3-none-any.whl

agentrail_python/_privacy.py

@@ -0,0 +1,255 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""Sanitized copies: what leaves the process, built from an allowlist.
+
+The customer's spans and log records are never changed; the queue only ever
+holds these copies.
+"""
+
+from __future__ import annotations
+
+import math
+import uuid
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit
+
+from opentelemetry._logs import LogRecord, SeverityNumber
+from opentelemetry.context import Context
+from opentelemetry.sdk._logs import ReadableLogRecord
+from opentelemetry.sdk.trace import ReadableSpan
+from opentelemetry.trace import (
+    INVALID_SPAN,
+    Link,
+    NonRecordingSpan,
+    SpanContext,
+    SpanKind,
+    Status,
+    TraceFlags,
+    set_span_in_context,
+)
+from opentelemetry.util.types import AttributeValue
+
+from ._handoff import valid_marker
+from ._runtime import Attributes, Runtime, SpanMeta
+from ._text import bounded, js_json, utf16_length, valid_text
+
+_LIMIT = 4096
+
+
+def clean_context(context: SpanContext) -> SpanContext:
+    """The IDs and flags only: trace state never leaves."""
+    return SpanContext(
+        trace_id=context.trace_id,
+        span_id=context.span_id,
+        is_remote=context.is_remote,
+        trace_flags=TraceFlags(context.trace_flags),
+    )
+
+
+def sanitized_url(raw: object) -> str | None:
+    """An http(s) URL with only valid ``h`` markers in its query; None past 4096 units."""
+    if not isinstance(raw, str):
+        return None
+    try:
+        parts = urlsplit(raw)
+        port = parts.port
+    except ValueError:
+        return None
+    if parts.scheme.lower() not in ("http", "https") or not parts.hostname:
+        return None
+    markers = [value for key, value in parse_qsl(parts.query, keep_blank_values=True) if key == "h"]
+    host = parts.hostname
+    if ":" in host:
+        host = f"[{host}]"
+    netloc = host if port is None else f"{host}:{port}"
+    query = urlencode([("h", marker) for marker in markers if valid_marker(marker)])
+    url = urlunsplit((parts.scheme.lower(), netloc, parts.path or "/", query, ""))
+    return url if valid_text(url, _LIMIT) else None
+
+
+def _first(attributes: Mapping[str, AttributeValue], *names: str) -> AttributeValue | None:
+    for name in names:
+        if name in attributes:
+            return attributes[name]
+    return None
+
+
+def copy_span(span: ReadableSpan, meta: SpanMeta, runtime: Runtime) -> ReadableSpan:
+    original: Mapping[str, AttributeValue] = span.attributes or {}
+    attributes: Attributes = {}
+    pairs = {
+        "http.request.method": ("http.request.method", "http.method"),
+        "http.route": ("http.route",),
+        "http.response.status_code": ("http.response.status_code", "http.status_code"),
+        "user_agent.original": ("user_agent.original", "http.user_agent"),
+        "error.type": ("error.type", "exception.type"),
+    }
+    for key, names in pairs.items():
+        value = _first(original, *names)
+        if isinstance(value, str):
+            attributes[key] = bounded(value, _LIMIT)
+        elif isinstance(value, int) and not isinstance(value, bool):
+            attributes[key] = value
+    if "error.type" not in attributes:
+        for event in span.events:
+            kind = (event.attributes or {}).get("exception.type")
+            if event.name == "exception" and isinstance(kind, str):
+                attributes["error.type"] = bounded(kind, _LIMIT)
+                break
+    address, path, scheme = original.get("server.address"), original.get("url.path"), original.get("url.scheme")
+    incoming = None
+    if isinstance(address, str) and isinstance(path, str) and isinstance(scheme, str):
+        port = original.get("server.port")
+        query = original.get("url.query")
+        incoming = f"{scheme}://{address}{f':{port}' if isinstance(port, int) else ''}{path}"
+        if isinstance(query, str):
+            incoming += f"?{query}"
+    url = sanitized_url(_first(original, "url.full", "http.url") or incoming)
+    if url is not None:
+        attributes["url.full"] = url
+
+    kind = span.kind
+    parent = span.parent
+    own = span.instrumentation_scope is not None and span.instrumentation_scope.name == runtime.scope.name
+    if kind in (SpanKind.SERVER, SpanKind.CONSUMER) and not own and parent is not None and not parent.is_remote:
+        # A library's own server span inside a managed operation (the MCP SDK opens
+        # one per message) is supporting work, not a second entry.
+        kind = SpanKind.INTERNAL
+    mark = meta.mark
+    entry = mark is not None or kind in (SpanKind.SERVER, SpanKind.CONSUMER)
+    if entry:
+        identity = mark.identity if mark is not None and mark.identity is not None else meta.state.principal()
+        attributes.update(identity)
+        attributes.update(meta.state.location)
+    if kind == SpanKind.SERVER:
+        attributes["agentrail.ingress.verified"] = meta.state.verified
+    method = attributes.get("http.request.method")
+    name = f"HTTP {method}" if isinstance(method, str) else "backend.operation"
+    if mark is not None:
+        name = mark.name
+        attributes["agentrail.step"] = True
+        attributes["agentrail.step.surface"] = mark.details.surface
+        attributes["agentrail.step.kind"] = mark.details.kind
+        if mark.performance is not None:
+            attributes.update(mark.performance)
+        mcp = mark.details.mcp
+        if mcp is not None:
+            for key, value, limit in (
+                ("agentrail.mcp.client.name", mcp.client, 128),
+                ("agentrail.mcp.client.version", mcp.client_version, 128),
+                ("agentrail.mcp.protocol.version", mcp.protocol, 64),
+            ):
+                if value:
+                    attributes[key] = bounded(value, limit)
+    links = [Link(clean_context(link.context)) for link in span.links]
+    if kind == SpanKind.SERVER:
+        links.extend(Link(clean_context(link.context)) for link in meta.state.links)
+    context = span.get_span_context()
+    assert context is not None  # an SDK span always has a context
+    return ReadableSpan(
+        name=bounded(name, _LIMIT),
+        context=clean_context(context),
+        parent=None if parent is None else clean_context(parent),
+        resource=runtime.resource,
+        attributes=attributes,
+        events=(),
+        links=links,
+        kind=kind,
+        status=Status(span.status.status_code),
+        start_time=span.start_time,
+        end_time=span.end_time,
+        instrumentation_scope=runtime.scope,
+    )
+
+
+def _scalar(value: object) -> str | int | float | bool | type[_Skip] | None:
+    if value is None or isinstance(value, bool):
+        return value
+    if isinstance(value, str):
+        return bounded(value, _LIMIT)
+    if isinstance(value, int):
+        return value
+    if isinstance(value, float):
+        return value if math.isfinite(value) else _Skip
+    return _Skip
+
+
+class _Skip:
+    """A field that is not a scalar and is left out."""
+
+
+def _fit(value: str, budget: int) -> str:
+    """The longest prefix whose JSON encoding fits ``budget`` UTF-16 units."""
+    low, high = 0, len(value)
+    while low < high:
+        middle = (low + high + 1) // 2
+        if utf16_length(js_json(value[:middle])) <= budget:
+            low = middle
+        else:
+            high = middle - 1
+    return value[:low]
+
+
+def selected_body(fields: Mapping[str, object], selected: tuple[str, ...]) -> str | None:
+    """A JSON object of the selected scalar fields, in allowlist order, at most 4096 units."""
+    result: dict[str, object] = {}
+    body: str | None = None
+    for key in selected:
+        if key not in fields or key in result:
+            continue
+        scalar = _scalar(fields[key])
+        if scalar is _Skip:
+            continue
+        candidate = {**result, key: scalar}
+        encoded = js_json(candidate)
+        if utf16_length(encoded) > _LIMIT:
+            if not isinstance(scalar, str):
+                continue
+            overhead = utf16_length(js_json({**result, key: ""})) - 2
+            budget = _LIMIT - overhead
+            if budget < 2:
+                continue
+            candidate[key] = _fit(scalar, budget)
+            encoded = js_json(candidate)
+        result = candidate
+        body = encoded
+    return body
+
+
+@dataclass(frozen=True, kw_only=True)
+class StructuredLog:
+    """A log entry. ``fields`` reach Agentrail only when named in ``logs.include_fields``."""
+
+    severity_number: int | None = None
+    severity_text: str | None = None
+    error_type: str | None = None
+    event_name: str | None = None
+    fields: Mapping[str, object] = field(default_factory=lambda: {})
+
+
+def copy_log(
+    entry: StructuredLog,
+    timestamp: int,
+    observed: int,
+    span: SpanContext | None,
+    runtime: Runtime,
+) -> ReadableLogRecord:
+    attributes: Attributes = {"agentrail.log.id": str(uuid.uuid4())}
+    if entry.error_type:
+        attributes["error.type"] = bounded(entry.error_type, _LIMIT)
+    body = selected_body(entry.fields, runtime.options.include_fields)
+    if body is not None:
+        attributes["agentrail.log.body_allowed"] = True
+    severity = entry.severity_number
+    record = LogRecord(
+        timestamp=timestamp,
+        observed_timestamp=observed,
+        # Exact correlation, never the ambient span: an uncorrelated record names none.
+        context=set_span_in_context(INVALID_SPAN if span is None else NonRecordingSpan(span), Context()),
+        severity_text=None if entry.severity_text is None else bounded(entry.severity_text, 64),
+        severity_number=SeverityNumber(severity) if severity is not None and 1 <= severity <= 24 else None,
+        body=body,
+        attributes=attributes,
+        event_name=None if not entry.event_name else bounded(entry.event_name, 256),
+    )
+    return ReadableLogRecord(record, resource=runtime.resource, instrumentation_scope=runtime.scope)

agentrail_python/_processors.py

@@ -0,0 +1,259 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""OpenTelemetry pieces a customer's own providers are built with."""
+
+from __future__ import annotations
+
+import os
+import time
+import weakref
+from collections.abc import Sequence
+from typing import cast
+
+import requests
+from opentelemetry import context as otel_context
+from opentelemetry import trace
+from opentelemetry.exporter.otlp.proto.http._log_exporter import OTLPLogExporter
+from opentelemetry.exporter.otlp.proto.http.trace_exporter import OTLPSpanExporter
+from opentelemetry.propagators.textmap import (
+    CarrierT,
+    Getter,
+    Setter,
+    TextMapPropagator,
+    default_getter,
+    default_setter,
+)
+from opentelemetry.sdk._logs import LogRecordProcessor, ReadableLogRecord, ReadWriteLogRecord
+from opentelemetry.sdk._logs.export import LogRecordExportResult
+from opentelemetry.sdk.trace import ReadableSpan, SpanProcessor
+from opentelemetry.sdk.trace import Span as SdkSpan
+from opentelemetry.sdk.trace.export import SpanExportResult
+from opentelemetry.sdk.trace.sampling import Decision, Sampler, SamplingResult
+from opentelemetry.trace import Link, SpanKind, TraceState
+from opentelemetry.util.types import Attributes
+
+from ._client import Client
+from ._handoff import OWNED_HEADERS
+from ._options import ConfigurationError, Options, resolve
+from ._privacy import StructuredLog, copy_log, copy_span
+from ._queue import ExportQueue
+from ._runtime import Runtime, SpanMeta
+from ._suppression import suppressed
+
+
+class _SpanProcessor(SpanProcessor):
+    def __init__(self, runtime: Runtime, queue: ExportQueue[ReadableSpan]) -> None:
+        self._runtime = runtime
+        self._queue = queue
+
+    def on_start(self, span: SdkSpan, parent_context: otel_context.Context | None = None) -> None:
+        if suppressed():
+            return
+        ctx = otel_context.get_current() if parent_context is None else parent_context
+        state = self._runtime.state(ctx)
+        if state is None:
+            return
+        context = span.get_span_context()
+        self._runtime.remember(context.trace_id, context.span_id, SpanMeta(state, self._runtime.mark(ctx)))
+
+    def on_end(self, span: ReadableSpan) -> None:
+        context = span.get_span_context()
+        if context is None:
+            return
+        meta = self._runtime.forget(context.trace_id, context.span_id)
+        if meta is not None:
+            self._queue.add(copy_span(span, meta, self._runtime))
+
+    def shutdown(self) -> None:
+        self._queue.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        self._queue.flush(timeout_millis / 1000)
+        return True
+
+
+class _LogProcessor(LogRecordProcessor):
+    """Copies an OTel log record's allowlisted fields; the record itself is untouched."""
+
+    def __init__(self, runtime: Runtime, queue: ExportQueue[ReadableLogRecord]) -> None:
+        self._runtime = runtime
+        self._queue = queue
+
+    def on_emit(self, log_record: ReadWriteLogRecord) -> None:
+        if suppressed():
+            return
+        if self._runtime.options.input.logs.exclude_uncorrelated and self._runtime.state() is None:
+            return
+        record = log_record.log_record
+        fields: dict[str, object] = {}
+        if isinstance(record.body, str):
+            fields["message"] = record.body
+        error_type: str | None = None
+        event_name = record.event_name
+        for key, value in (record.attributes or {}).items():
+            fields[key] = value
+            if key in ("error.type", "exception.type") and isinstance(value, str):
+                error_type = value
+            if key == "event.name" and isinstance(value, str):
+                event_name = value
+        span = None
+        if record.trace_id and record.span_id:
+            span = trace.SpanContext(record.trace_id, record.span_id, False, trace.TraceFlags(record.trace_flags or 0))
+        now = time.time_ns()
+        entry = StructuredLog(
+            severity_number=None if record.severity_number is None else record.severity_number.value,
+            severity_text=record.severity_text,
+            error_type=error_type,
+            event_name=event_name,
+            fields=fields,
+        )
+        self._queue.add(copy_log(entry, record.timestamp or now, record.observed_timestamp or now, span, self._runtime))
+
+    def shutdown(self) -> None:
+        self._queue.shutdown()
+
+    def force_flush(self, timeout_millis: int = 30000) -> bool:
+        self._queue.flush(timeout_millis / 1000)
+        return True
+
+
+class _Sampler(Sampler):
+    """Promotes a dropped span inside a managed operation to record-only."""
+
+    def __init__(self, runtime: Runtime, base: Sampler) -> None:
+        self._runtime = runtime
+        self._base = base
+
+    def should_sample(
+        self,
+        parent_context: otel_context.Context | None,
+        trace_id: int,
+        name: str,
+        kind: SpanKind | None = None,
+        attributes: Attributes = None,
+        links: Sequence[Link] | None = None,
+        trace_state: TraceState | None = None,
+    ) -> SamplingResult:
+        result = self._base.should_sample(parent_context, trace_id, name, kind, attributes, links, trace_state)
+        ctx = otel_context.get_current() if parent_context is None else parent_context
+        if result.decision == Decision.DROP and not suppressed() and self._runtime.state(ctx) is not None:
+            # The SDK starts a span with the sampler's attributes only, and a DROP
+            # carries none: keep the ones the span was started with.
+            returned = cast(Attributes, result.attributes)  # pyright: ignore[reportUnknownMemberType]
+            merged = {**(attributes or {}), **(returned or {})}
+            return SamplingResult(Decision.RECORD_ONLY, merged, result.trace_state)
+        return result
+
+    def get_description(self) -> str:
+        return f"Agentrail({self._base.get_description()})"
+
+
+class _Propagator(TextMapPropagator):
+    """Managed work is never continued from, or injected with, an unsigned header.
+
+    Unmanaged work goes to the customer's propagator.
+    """
+
+    def __init__(self, runtime: Runtime, base: TextMapPropagator | None) -> None:
+        self._runtime = runtime
+        self._base = base
+
+    def extract(
+        self,
+        carrier: CarrierT,
+        context: otel_context.Context | None = None,
+        getter: Getter[CarrierT] = default_getter,  # type: ignore[assignment]
+    ) -> otel_context.Context:
+        ctx = otel_context.get_current() if context is None else context
+        if self._runtime.state(ctx) is not None or self._base is None:
+            return ctx
+        return self._base.extract(carrier, context, getter)
+
+    def inject(
+        self,
+        carrier: CarrierT,
+        context: otel_context.Context | None = None,
+        setter: Setter[CarrierT] = default_setter,  # type: ignore[assignment]
+    ) -> None:
+        ctx = otel_context.get_current() if context is None else context
+        if self._runtime.state(ctx) is not None or suppressed() or self._base is None:
+            return
+        self._base.inject(carrier, context, setter)
+
+    @property
+    def fields(self) -> set[str]:
+        return set(OWNED_HEADERS) | (set() if self._base is None else set(self._base.fields))
+
+
+def _export_spans(exporter: OTLPSpanExporter, spans: Sequence[ReadableSpan]) -> bool:
+    return exporter.export(spans) == SpanExportResult.SUCCESS
+
+
+def _export_logs(exporter: OTLPLogExporter, records: Sequence[ReadableLogRecord]) -> bool:
+    return exporter.export(records) == LogRecordExportResult.SUCCESS
+
+
+class Processors:
+    """Construct before the customer's tracer provider; then ``bind`` that provider once.
+
+    Registers nothing globally.
+    """
+
+    def __init__(self, options: Options) -> None:
+        resolved = resolve(options)
+        self._runtime = runtime = Runtime(resolved)
+        headers = {"Authorization": f"Bearer {options.collector.secret}"}
+        timeout = resolved.queue.export_timeout
+        span_session, log_session = requests.Session(), requests.Session()
+        span_exporter = OTLPSpanExporter(
+            endpoint=f"{resolved.endpoint}/v1/traces", headers=headers, timeout=timeout, session=span_session
+        )
+        log_exporter = OTLPLogExporter(
+            endpoint=f"{resolved.endpoint}/v1/logs", headers=headers, timeout=timeout, session=log_session
+        )
+        # A forked child must not write to a connection the parent still uses:
+        # it drops the pooled connections it inherited and opens its own.
+        # Weak references: a shut-down SDK's sessions are not kept alive by the hook.
+        if hasattr(os, "register_at_fork"):
+            sessions = (weakref.ref(span_session), weakref.ref(log_session))
+
+            def drop_inherited_connections() -> None:
+                for reference in sessions:
+                    session = reference()
+                    if session is not None:
+                        session.close()
+
+            os.register_at_fork(after_in_child=drop_inherited_connections)
+        self.spans: ExportQueue[ReadableSpan] = ExportQueue(
+            "traces",
+            resolved.queue,
+            lambda batch: _export_spans(span_exporter, batch),
+            lambda code, count: runtime.report(code, "traces", count),
+            span_exporter.shutdown,
+        )
+        self.logs: ExportQueue[ReadableLogRecord] = ExportQueue(
+            "logs",
+            resolved.queue,
+            lambda batch: _export_logs(log_exporter, batch),
+            lambda code, count: runtime.report(code, "logs", count),
+            log_exporter.shutdown,
+        )
+        self.span_processor: SpanProcessor = _SpanProcessor(runtime, self.spans)
+        self.log_processor: LogRecordProcessor = _LogProcessor(runtime, self.logs)
+        self.propagator: TextMapPropagator = _Propagator(runtime, options.propagator)
+        self._bound = False
+
+    def sampler(self, base: Sampler) -> Sampler:
+        """Wrap the customer's sampler so managed spans it drops stay record-only."""
+        return _Sampler(self._runtime, base)
+
+    def bind(self, provider: trace.TracerProvider) -> Client:
+        """The client that records with ``provider``. Call once."""
+        if self._bound:
+            raise ConfigurationError("These processors are already bound to a tracer provider.")
+        self._bound = True
+        return Client(self._runtime, provider, self.spans, self.logs)
+
+
+def create_processors(options: Options) -> Processors:
+    """Agentrail's sampler wrapper, span and log processors and propagator for an existing setup."""
+    return Processors(options)