agentrail-python 0.1.0__py3-none-any.whl

agentrail_python/_queue.py

@@ -0,0 +1,236 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""Bounded, best-effort batch delivery on one background thread per signal.
+
+Capacity counts the batch being exported. One batch is exported at a time; the
+OTLP exporter owns retries within its timeout. Records only ever wait in
+memory: there is no disk persistence.
+"""
+
+from __future__ import annotations
+
+import os
+import threading
+import time
+import weakref
+from collections import deque
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from typing import Generic, TypeVar
+
+from ._options import CaptureCode, QueueOptions
+from ._suppression import reporting
+
+T = TypeVar("T")
+
+
+@dataclass(frozen=True)
+class QueueStats:
+    accepted: int
+    exported: int
+    dropped: int
+    pending: int
+
+
+class _Batch(Generic[T]):
+    def __init__(self, records: list[tuple[int, T]]) -> None:
+        self.records = records
+        self.dropped = False
+
+
+class ExportQueue(Generic[T]):
+    def __init__(
+        self,
+        name: str,
+        options: QueueOptions,
+        export: Callable[[Sequence[T]], bool],
+        report: Callable[[CaptureCode, int], None],
+        shutdown_exporter: Callable[[], None],
+    ) -> None:
+        self._options = options
+        self._export = export
+        self._report = report
+        self._shutdown_exporter = shutdown_exporter
+        self._lock = threading.Condition()
+        self._records: deque[tuple[int, T]] = deque()
+        self._active: _Batch[T] | None = None
+        self._accepted = 0
+        self._exported = 0
+        self._dropped = 0
+        self._pending = 0
+        self._force = 0
+        self._closed = False
+        self._stop_once = threading.Lock()
+        self._stopped: threading.Event | None = None
+        self._name = name
+        self._worker = threading.Thread(target=self._run, name=f"agentrail-{name}", daemon=True)
+        self._worker.start()
+        # A process that forks after start-up (gunicorn --preload, multiprocessing)
+        # gives the child this queue without its worker thread, and possibly with
+        # a lock held by a thread that no longer exists. The child starts afresh.
+        if hasattr(os, "register_at_fork"):
+            queue = weakref.ref(self)
+
+            def reinitialize() -> None:
+                alive = queue()
+                if alive is not None:
+                    alive._after_fork()  # pyright: ignore[reportPrivateUsage]
+
+            os.register_at_fork(after_in_child=reinitialize)
+
+    def _after_fork(self) -> None:
+        """In a forked child: the parent's records stay the parent's to send."""
+        self._lock = threading.Condition()
+        self._records = deque()
+        self._active = None
+        self._accepted = 0
+        self._exported = 0
+        self._dropped = 0
+        self._pending = 0
+        self._force = 0
+        self._stop_once = threading.Lock()
+        if self._closed:
+            return
+        self._stopped = None
+        self._worker = threading.Thread(target=self._run, name=f"agentrail-{self._name}", daemon=True)
+        self._worker.start()
+
+    def diagnostics(self) -> QueueStats:
+        with self._lock:
+            return QueueStats(self._accepted, self._exported, self._dropped, self._pending)
+
+    def add(self, value: T) -> bool:
+        """Never waits for an export. False means the record was dropped."""
+        with self._lock:
+            retained = len(self._records) + (0 if self._active is None else len(self._active.records))
+            code: CaptureCode | None = None
+            if self._closed:
+                code = "closed"
+            elif retained >= self._options.max_records:
+                code = "queue_full"
+            if code is None:
+                self._accepted += 1
+                self._pending += 1
+                self._records.append((self._accepted, value))
+                if len(self._records) >= self._options.batch_size:
+                    self._lock.notify_all()
+                return True
+            self._dropped += 1
+        self._report(code, 1)
+        return False
+
+    def flush(self, timeout: float | None = None) -> None:
+        """Attempts every record accepted when called. Delivery is not guaranteed."""
+        limit = self._options.flush_timeout if timeout is None else min(timeout, self._options.flush_timeout)
+        with self._lock:
+            target = self._accepted
+            self._force = max(self._force, target)
+            self._lock.notify_all()
+        self._wait(target, time.monotonic() + limit)
+
+    def shutdown(self, timeout: float | None = None) -> None:
+        """Closes admission, drains once within the shutdown deadline, then stops."""
+        with self._stop_once:
+            first = self._stopped is None
+            if first:
+                self._stopped = threading.Event()
+        stopped = self._stopped
+        assert stopped is not None
+        if not first:
+            stopped.wait(self._options.shutdown_timeout)
+            return
+        limit = self._options.shutdown_timeout if timeout is None else min(timeout, self._options.shutdown_timeout)
+        deadline = time.monotonic() + limit
+        with self._lock:
+            self._closed = True
+            target = self._accepted
+            self._force = target
+            self._lock.notify_all()
+        self._wait(target, deadline)
+        self._worker.join(max(0.0, deadline - time.monotonic()))
+        with reporting():
+            self._shutdown_exporter()
+        stopped.set()
+
+    def _outstanding(self, target: int) -> bool:
+        active = self._active
+        if active is not None and not active.dropped and active.records[0][0] <= target:
+            return True
+        return len(self._records) > 0 and self._records[0][0] <= target
+
+    def _wait(self, target: int, deadline: float) -> None:
+        count = 0
+        with self._lock:
+            while self._outstanding(target):
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    count = self._abort(target)
+                    break
+                self._lock.wait(remaining)
+        if count > 0:
+            self._report("deadline", count)
+
+    def _abort(self, target: int) -> int:
+        count = 0
+        active = self._active
+        if active is not None and not active.dropped and active.records[0][0] <= target:
+            active.dropped = True
+            count += len(active.records)
+        while self._records and self._records[0][0] <= target:
+            self._records.popleft()
+            count += 1
+        self._dropped += count
+        self._pending -= count
+        self._lock.notify_all()
+        return count
+
+    def _next_batch(self) -> _Batch[T] | None:
+        """Called with the lock held. None when the worker should wait or stop."""
+        size = self._options.batch_size
+        records = self._records
+        if not records:
+            return None
+        if not self._closed and len(records) < size and records[0][0] > self._force:
+            return None
+        count = min(len(records), size)
+        # A partial batch never sweeps in records accepted after the flush that forced it.
+        if not self._closed and records[0][0] <= self._force:
+            while count > 1 and records[count - 1][0] > self._force:
+                count -= 1
+        batch = _Batch([records.popleft() for _ in range(count)])
+        self._active = batch
+        return batch
+
+    def _run(self) -> None:
+        with reporting():
+            next_tick = time.monotonic() + self._options.scheduled_delay
+            while True:
+                with self._lock:
+                    batch = self._next_batch()
+                    while batch is None:
+                        if self._closed and not self._records:
+                            return
+                        remaining = next_tick - time.monotonic()
+                        if remaining <= 0:
+                            self._force = self._accepted
+                            next_tick = time.monotonic() + self._options.scheduled_delay
+                        else:
+                            self._lock.wait(remaining)
+                        batch = self._next_batch()
+                values = [value for _, value in batch.records]
+                try:
+                    succeeded = self._export(values)
+                except Exception:  # noqa: BLE001 - a failed export is reported as loss, never raised
+                    succeeded = False
+                with self._lock:
+                    count = len(batch.records)
+                    reported = not batch.dropped
+                    if reported:
+                        self._pending -= count
+                        if succeeded:
+                            self._exported += count
+                        else:
+                            self._dropped += count
+                    self._active = None
+                    self._lock.notify_all()
+                if reported and not succeeded:
+                    self._report("export_failed", count)

agentrail_python/_runtime.py

@@ -0,0 +1,110 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""Per-installation state: context keys, the resource, and span metadata."""
+
+from __future__ import annotations
+
+import threading
+import uuid
+from dataclasses import dataclass, field
+
+from opentelemetry import context as otel_context
+from opentelemetry.sdk.resources import Resource
+from opentelemetry.sdk.util.instrumentation import InstrumentationScope
+from opentelemetry.trace import Link
+from opentelemetry.util.types import AttributeValue
+
+from ._handoff import key_id, keyed_hash
+from ._options import CaptureCode, CaptureError, MilestoneContext, Resolved, Signal, VerifiedIdentity
+from ._version import VERSION
+
+COMPONENT = "agentrail-python"
+Attributes = dict[str, AttributeValue]
+
+
+@dataclass
+class RequestState:
+    """One managed operation's evidence, shared by its child spans and tasks."""
+
+    verified: bool = False
+    links: list[Link] = field(default_factory=lambda: [])
+    location: Attributes = field(default_factory=lambda: {})
+    _identity: Attributes | None = None
+    _lock: threading.Lock = field(default_factory=threading.Lock)
+
+    def principal(self) -> Attributes:
+        with self._lock:
+            return dict(self._identity or {})
+
+    def assert_identity(self, identity: Attributes) -> bool:
+        """Keeps the first assertion. True when ``identity`` conflicts with it."""
+        with self._lock:
+            if self._identity is None:
+                self._identity = identity
+                return False
+            return self._identity != identity
+
+
+@dataclass(frozen=True)
+class Mark:
+    name: str
+    details: MilestoneContext
+    performance: Attributes | None = None
+    identity: Attributes | None = None
+
+
+@dataclass(frozen=True)
+class SpanMeta:
+    state: RequestState
+    mark: Mark | None
+
+
+class Runtime:
+    def __init__(self, options: Resolved) -> None:
+        self.options = options
+        unique = uuid.uuid4().hex
+        self.state_key = otel_context.create_key(f"agentrail.state.{unique}")
+        self.mark_key = otel_context.create_key(f"agentrail.milestone.{unique}")
+        self.key_id = key_id(options.input.hash_key)
+        self.resource = Resource(
+            {
+                "service.name": options.input.service_name,
+                "service.instance.id": options.instance,
+                "agentrail.component.name": COMPONENT,
+                "agentrail.component.version": VERSION,
+                "agentrail.hash_key_id": self.key_id,
+                "agentrail.capture.surfaces": "http,mcp,auth,custom",
+            }
+        )
+        self.scope = InstrumentationScope(COMPONENT, VERSION)
+        self._spans: dict[tuple[int, int], SpanMeta] = {}
+        self._spans_lock = threading.Lock()
+
+    def state(self, ctx: otel_context.Context | None = None) -> RequestState | None:
+        value = otel_context.get_value(self.state_key, ctx)
+        return value if isinstance(value, RequestState) else None
+
+    def mark(self, ctx: otel_context.Context | None = None) -> Mark | None:
+        value = otel_context.get_value(self.mark_key, ctx)
+        return value if isinstance(value, Mark) else None
+
+    def remember(self, trace_id: int, span_id: int, meta: SpanMeta) -> None:
+        with self._spans_lock:
+            self._spans[(trace_id, span_id)] = meta
+
+    def forget(self, trace_id: int, span_id: int) -> SpanMeta | None:
+        with self._spans_lock:
+            return self._spans.pop((trace_id, span_id), None)
+
+    def identity(self, identity: VerifiedIdentity) -> Attributes:
+        key = self.options.input.hash_key
+        attributes: Attributes = {
+            "agentrail.identity.verified": True,
+            "agentrail.principal.kind": identity.kind,
+            "agentrail.principal.id_hash": keyed_hash(key, identity.id),
+        }
+        if identity.acting_user_id is not None:
+            attributes["agentrail.principal.acting_user_hash"] = keyed_hash(key, identity.acting_user_id)
+        return attributes
+
+    def report(self, code: CaptureCode, signal: Signal, count: int = 1) -> None:
+        self.options.report(CaptureError(code, signal, count))

agentrail_python/_suppression.py

@@ -0,0 +1,33 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""Work the SDK does for itself is never captured.
+
+Error reporting and export run under this flag so a callback that logs, or an
+exporter that logs its own failures, cannot feed the queue it reports on.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Generator
+from contextlib import contextmanager
+from contextvars import ContextVar
+
+from opentelemetry import context as otel_context
+
+_flag: ContextVar[bool] = ContextVar("agentrail_python_suppressed", default=False)
+_OTEL_KEY = otel_context._SUPPRESS_INSTRUMENTATION_KEY  # pyright: ignore[reportPrivateUsage]
+
+
+def suppressed() -> bool:
+    return _flag.get()
+
+
+@contextmanager
+def reporting() -> Generator[None]:
+    """Run SDK-internal work with capture and OTel instrumentation suppressed."""
+    flag = _flag.set(True)
+    token = otel_context.attach(otel_context.set_value(_OTEL_KEY, True))
+    try:
+        yield
+    finally:
+        otel_context.detach(token)
+        _flag.reset(flag)

agentrail_python/_text.py

@@ -0,0 +1,49 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""Text limits counted the way the collector and the JavaScript SDKs count them.
+
+Limits are UTF-16 code units, because the collector and agentrail-node measure
+JavaScript string length. JSON is encoded exactly as ``JSON.stringify`` does for
+well-formed strings, so hashes and signatures agree across the SDKs.
+"""
+
+from __future__ import annotations
+
+import json
+
+
+def utf16_length(value: str) -> int:
+    return sum(2 if ord(character) > 0xFFFF else 1 for character in value)
+
+
+def well_formed(value: str) -> bool:
+    """True when the string has no lone surrogate, so it encodes as UTF-8."""
+    try:
+        value.encode("utf-8")
+    except UnicodeEncodeError:
+        return False
+    return True
+
+
+def valid_text(value: str, limit: int) -> bool:
+    """Well-formed, without NUL, and at most ``limit`` UTF-16 units."""
+    return well_formed(value) and "\x00" not in value and utf16_length(value) <= limit
+
+
+def bounded(value: str, limit: int) -> str:
+    """Drop NUL, replace lone surrogates and cut to ``limit`` UTF-16 units.
+
+    A cut never splits a code point.
+    """
+    cleaned = value.encode("utf-8", "replace").decode("utf-8").replace("\x00", "")
+    units = 0
+    for index, character in enumerate(cleaned):
+        width = 2 if ord(character) > 0xFFFF else 1
+        if units + width > limit:
+            return cleaned[:index]
+        units += width
+    return cleaned
+
+
+def js_json(value: object) -> str:
+    """``JSON.stringify`` for strings, finite numbers, booleans, null and lists of them."""
+    return json.dumps(value, ensure_ascii=False, separators=(",", ":"), allow_nan=False)

agentrail_python/_version.py

@@ -0,0 +1,2 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+VERSION = "0.1.0"

agentrail_python/asgi.py

@@ -0,0 +1,111 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""ASGI middleware: FastAPI, Starlette, Django (ASGI), Quart, Litestar."""
+
+from __future__ import annotations
+
+from collections.abc import Awaitable, Callable, MutableMapping
+from typing import TYPE_CHECKING, Any
+from urllib.parse import quote
+
+from opentelemetry import context as otel_context
+from opentelemetry.trace import StatusCode
+
+from ._options import IncomingRequest
+
+if TYPE_CHECKING:
+    from ._client import Client
+
+Scope = MutableMapping[str, Any]
+Message = MutableMapping[str, Any]
+Receive = Callable[[], Awaitable[Message]]
+Send = Callable[[Message], Awaitable[None]]
+ASGIApp = Callable[[Scope, Receive, Send], Awaitable[None]]
+
+
+def _request(scope: Scope) -> IncomingRequest:
+    headers: dict[str, list[str]] = {}
+    for raw_name, raw_value in scope.get("headers", ()):
+        name = bytes(raw_name).decode("latin-1").lower()
+        headers.setdefault(name, []).append(bytes(raw_value).decode("latin-1"))
+    raw_path = scope.get("raw_path")
+    path = (
+        bytes(raw_path).decode("latin-1")
+        if raw_path
+        else quote(str(scope.get("root_path", "")) + str(scope.get("path", "/")), safe="/%:@!$&'()*+,;=~")
+    )
+    query = bytes(scope.get("query_string", b"")).decode("latin-1")
+    target = f"{path}?{query}" if query else path
+    host = headers.get("host", [""])[0]
+    if not host:
+        server = scope.get("server")
+        host = f"{server[0]}:{server[1]}" if server else "localhost"
+    client = scope.get("client")
+    return IncomingRequest(
+        method=str(scope.get("method", "GET")),
+        target=target,
+        scheme=str(scope.get("scheme", "http")),
+        host=host,
+        # A header sent more than once is joined, so a handoff or traceparent
+        # sent twice cannot verify.
+        headers={name: ", ".join(values) for name, values in headers.items()},
+        client=None if not client else str(client[0]),
+        raw=scope,
+    )
+
+
+def _route(scope: Scope) -> str | None:
+    route = scope.get("route")
+    for attribute in ("path_format", "path"):
+        value = getattr(route, attribute, None)
+        if isinstance(value, str) and value:
+            return value
+    return None
+
+
+class AgentrailASGIMiddleware:
+    """Records each HTTP request the wrapped application serves as a journey step.
+
+    ``app.add_middleware(AgentrailASGIMiddleware, client=agentrail)`` in FastAPI and
+    Starlette, or ``application = agentrail.asgi(application)``. Install it once.
+    """
+
+    def __init__(self, app: ASGIApp, *, client: Client) -> None:
+        self.app = app
+        self._client = client
+
+    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
+        if scope.get("type") != "http":
+            await self.app(scope, receive, send)
+            return
+        entered = self._client._enter(_request(scope))  # pyright: ignore[reportPrivateUsage]
+        if entered is None:
+            await self.app(scope, receive, send)
+            return
+        span, ctx = entered
+        status: list[int] = []
+
+        async def observed_send(message: Message) -> None:
+            if message.get("type") == "http.response.start" and not status:
+                status.append(int(message["status"]))
+            await send(message)
+
+        token = otel_context.attach(ctx)
+        try:
+            await self.app(scope, receive, observed_send)
+        except BaseException as error:
+            span.set_status(StatusCode.ERROR)
+            span.set_attribute("error.type", type(error).__qualname__)
+            raise
+        finally:
+            try:
+                code = status[0] if status else None
+                if code is not None:
+                    span.set_attribute("http.response.status_code", code)
+                    if code >= 500:
+                        span.set_status(StatusCode.ERROR)
+                route = _route(scope)
+                if route is not None:
+                    span.set_attribute("http.route", route)
+                span.end()
+            finally:
+                otel_context.detach(token)