agentrail-python 0.1.0__py3-none-any.whl

agentrail_python/httpx.py

@@ -0,0 +1,159 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""httpx transports that sign trace context only for allowed origins.
+
+``httpx.Client(transport=agentrail_python.httpx.transport(agentrail))``. Each
+redirect passes through the transport again, so every actual destination is
+admitted on its own.
+"""
+
+from __future__ import annotations
+
+from collections.abc import AsyncIterator, Iterator
+from typing import TYPE_CHECKING
+
+import httpx
+
+from ._handoff import OWNED_HEADERS
+from ._outbound import begin, finish, origin
+
+if TYPE_CHECKING:
+    from opentelemetry.trace import Span
+
+    from ._client import Client
+
+
+def _prepare(client: Client, request: httpx.Request) -> Span | None:
+    url = request.url
+    target = url.raw_path.decode("ascii")
+    outgoing = begin(
+        client,
+        request.method,
+        origin(url.scheme, url.host, url.port),
+        target.split("?", 1)[0],
+        target,
+        str(url),
+    )
+    if outgoing.strip:
+        for name in OWNED_HEADERS:
+            request.headers.pop(name, None)
+    for name, value in outgoing.headers.items():
+        request.headers[name] = value
+    return outgoing.span
+
+
+class _Stream(httpx.SyncByteStream):
+    def __init__(self, stream: httpx.SyncByteStream, span: Span, status: int) -> None:
+        self._stream = stream
+        self._span = span
+        self._status = status
+        self._done = False
+
+    def _finish(self, error: BaseException | None) -> None:
+        if not self._done:
+            self._done = True
+            finish(self._span, self._status, error)
+
+    def __iter__(self) -> Iterator[bytes]:
+        try:
+            yield from self._stream
+        except BaseException as error:
+            self._finish(error)
+            raise
+
+    def close(self) -> None:
+        try:
+            self._stream.close()
+        except BaseException as error:
+            self._finish(error)
+            raise
+        self._finish(None)
+
+
+class _AsyncStream(httpx.AsyncByteStream):
+    def __init__(self, stream: httpx.AsyncByteStream, span: Span, status: int) -> None:
+        self._stream = stream
+        self._span = span
+        self._status = status
+        self._done = False
+
+    def _finish(self, error: BaseException | None) -> None:
+        if not self._done:
+            self._done = True
+            finish(self._span, self._status, error)
+
+    async def __aiter__(self) -> AsyncIterator[bytes]:
+        try:
+            async for chunk in self._stream:
+                yield chunk
+        except BaseException as error:
+            self._finish(error)
+            raise
+
+    async def aclose(self) -> None:
+        try:
+            await self._stream.aclose()
+        except BaseException as error:
+            self._finish(error)
+            raise
+        self._finish(None)
+
+
+class AgentrailTransport(httpx.BaseTransport):
+    def __init__(self, client: Client, base: httpx.BaseTransport | None = None) -> None:
+        self._client = client
+        self._base = base if base is not None else httpx.HTTPTransport()
+
+    def handle_request(self, request: httpx.Request) -> httpx.Response:
+        span = _prepare(self._client, request)
+        if span is None:
+            return self._base.handle_request(request)
+        try:
+            response = self._base.handle_request(request)
+        except BaseException as error:
+            finish(span, None, error)
+            raise
+        stream = response.stream
+        if response.is_closed or not isinstance(stream, httpx.SyncByteStream):
+            # A body already read (or none to read) ends the call now.
+            finish(span, response.status_code, None)
+            return response
+        response.stream = _Stream(stream, span, response.status_code)
+        return response
+
+    def close(self) -> None:
+        self._base.close()
+
+
+class AgentrailAsyncTransport(httpx.AsyncBaseTransport):
+    def __init__(self, client: Client, base: httpx.AsyncBaseTransport | None = None) -> None:
+        self._client = client
+        self._base = base if base is not None else httpx.AsyncHTTPTransport()
+
+    async def handle_async_request(self, request: httpx.Request) -> httpx.Response:
+        span = _prepare(self._client, request)
+        if span is None:
+            return await self._base.handle_async_request(request)
+        try:
+            response = await self._base.handle_async_request(request)
+        except BaseException as error:
+            finish(span, None, error)
+            raise
+        stream = response.stream
+        if response.is_closed or not isinstance(stream, httpx.AsyncByteStream):
+            finish(span, response.status_code, None)
+            return response
+        response.stream = _AsyncStream(stream, span, response.status_code)
+        return response
+
+    async def aclose(self) -> None:
+        await self._base.aclose()
+
+
+def transport(client: Client, base: httpx.BaseTransport | None = None) -> AgentrailTransport:
+    """A transport for ``httpx.Client``. ``base`` defaults to ``httpx.HTTPTransport()``."""
+    return AgentrailTransport(client, base)
+
+
+def async_transport(client: Client, base: httpx.AsyncBaseTransport | None = None) -> AgentrailAsyncTransport:
+    """A transport for ``httpx.AsyncClient``. ``base`` defaults to ``httpx.AsyncHTTPTransport()``."""
+    return AgentrailAsyncTransport(client, base)

agentrail_python/logging.py

@@ -0,0 +1,70 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""A ``logging.Handler`` that records structured logs, correlated with the request.
+
+``logging.getLogger().addHandler(AgentrailHandler(agentrail))``. The record is
+captured where it is emitted, so it carries that request's trace. Other
+handlers' output is unchanged.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import TYPE_CHECKING
+
+from ._privacy import StructuredLog
+
+if TYPE_CHECKING:
+    from ._client import Client
+
+# Attributes every LogRecord has; anything else came from ``extra=``.
+_STANDARD = frozenset(vars(logging.LogRecord("", 0, "", 0, "", None, None)).keys() | {"message", "asctime", "taskName"})
+
+
+def severity_number(level: int) -> int:
+    """OpenTelemetry's mapping of a logging level to a severity number."""
+    if level >= logging.CRITICAL:
+        return 21
+    if level >= logging.ERROR:
+        return 17
+    if level >= logging.WARNING:
+        return 13
+    if level >= logging.INFO:
+        return 9
+    if level >= logging.DEBUG:
+        return 5
+    return 1
+
+
+class AgentrailHandler(logging.Handler):
+    """Message text and ``extra`` fields leave only when named in ``logs.include_fields``.
+
+    ``message`` is the formatted message. ``event.name`` (or ``event_name``) in
+    ``extra`` names the event; an exception's class name becomes ``error.type``.
+    """
+
+    def __init__(self, client: Client, level: int = logging.NOTSET) -> None:
+        super().__init__(level)
+        self._client = client
+
+    def emit(self, record: logging.LogRecord) -> None:
+        try:
+            fields: dict[str, object] = {"message": record.getMessage()}
+            for key, value in vars(record).items():
+                if key not in _STANDARD:
+                    fields[key] = value
+            event = fields.get("event.name", fields.get("event_name"))
+            error = record.exc_info[1] if record.exc_info else None
+            explicit = fields.get("error.type")
+            self._client.log(
+                StructuredLog(
+                    severity_number=severity_number(record.levelno),
+                    severity_text=record.levelname,
+                    error_type=explicit
+                    if isinstance(explicit, str)
+                    else (type(error).__qualname__ if error is not None else None),
+                    event_name=event if isinstance(event, str) else None,
+                    fields=fields,
+                )
+            )
+        except Exception:  # noqa: BLE001 - logging's own contract: report through handleError
+            self.handleError(record)

agentrail_python/mcp.py

@@ -0,0 +1,125 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""MCP steps for the official ``mcp`` Python SDK (2.x, ``MCPServer``).
+
+Wrap the functions you register with a stable registered name; add the
+middleware so each step knows the calling client and protocol::
+
+    capture = mcp_capture(agentrail)
+    server = MCPServer("orders", middleware=[capture.middleware])
+    server.tool(name="checkout")(capture.tool("checkout", checkout))
+
+Arguments, results, prompt text and resource URIs are never recorded. Client
+details come from the client and are descriptive only: they never assert an
+identity.
+"""
+
+from __future__ import annotations
+
+import contextvars
+import functools
+import inspect
+from collections.abc import Awaitable, Callable, Mapping
+from typing import TYPE_CHECKING, Any, TypeVar, cast
+
+from opentelemetry import trace
+from opentelemetry.trace import StatusCode
+
+from ._options import McpContext, MilestoneContext
+from ._text import bounded
+
+if TYPE_CHECKING:
+    from ._client import Client
+
+F = TypeVar("F", bound=Callable[..., Any])
+
+_request: contextvars.ContextVar[tuple[str | None, str | None, str | None] | None] = contextvars.ContextVar(
+    "agentrail_python_mcp_request", default=None
+)
+
+
+def _text(value: object, limit: int) -> str | None:
+    return bounded(value, limit) or None if isinstance(value, str) else None
+
+
+def _span_name(method: str, name: str) -> str:
+    return bounded(f"mcp.{method}:{name}", 256)
+
+
+def _is_error(result: object) -> bool:
+    if isinstance(result, Mapping):
+        return cast(Mapping[str, object], result).get("isError") is True
+    return getattr(result, "is_error", None) is True or getattr(result, "isError", None) is True
+
+
+class McpCapture:
+    def __init__(self, client: Client) -> None:
+        self._client = client
+
+    async def middleware(self, ctx: Any, call_next: Callable[[Any], Awaitable[Any]]) -> Any:
+        """``ServerMiddleware``: keeps this request's client name, version and protocol."""
+        info = getattr(getattr(getattr(ctx, "session", None), "client_params", None), "client_info", None)
+        token = _request.set(
+            (
+                _text(getattr(info, "name", None), 128),
+                _text(getattr(info, "version", None), 128),
+                _text(getattr(ctx, "protocol_version", None), 64),
+            )
+        )
+        try:
+            return await call_next(ctx)
+        finally:
+            _request.reset(token)
+
+    def tool(self, name: str, handler: F) -> F:
+        """Wrap a tool function. A result with ``isError`` marks the step failed."""
+        return self._wrap("tools/call", "tool.call", name, handler)
+
+    def resource(self, name: str, handler: F) -> F:
+        """Wrap a resource function. Pass the registered name, never the requested URI."""
+        return self._wrap("resources/read", "mcp.resources/read", name, handler)
+
+    def prompt(self, name: str, handler: F) -> F:
+        """Wrap a prompt function."""
+        return self._wrap("prompts/get", "mcp.prompts/get", name, handler)
+
+    def _details(self, method: str, kind: str) -> MilestoneContext:
+        client, version, protocol = _request.get() or (None, None, None)
+        return MilestoneContext(
+            surface="mcp",
+            kind=kind,
+            mcp=McpContext(method=method, client=client, client_version=version, protocol=protocol),
+        )
+
+    def _wrap(self, method: str, kind: str, name: str, handler: F) -> F:
+        label = _span_name(method, name)
+        client = self._client
+
+        def mark_result(result: object) -> None:
+            if _is_error(result):
+                span = trace.get_current_span()
+                span.set_status(StatusCode.ERROR)
+                span.set_attribute("error.type", "McpOperationError")
+
+        if inspect.iscoroutinefunction(handler):
+
+            @functools.wraps(handler)
+            async def run_async(*args: Any, **kwargs: Any) -> Any:
+                async with client.start_milestone(label, self._details(method, kind)):
+                    result = await handler(*args, **kwargs)
+                    mark_result(result)
+                    return result
+
+            return cast(F, run_async)
+
+        @functools.wraps(handler)
+        def run(*args: Any, **kwargs: Any) -> Any:
+            with client.start_milestone(label, self._details(method, kind)):
+                result = handler(*args, **kwargs)
+                mark_result(result)
+                return result
+
+        return cast(F, run)
+
+
+def mcp_capture(client: Client) -> McpCapture:
+    return McpCapture(client)

agentrail_python/requests.py

@@ -0,0 +1,67 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""A requests adapter that signs trace context only for allowed origins.
+
+``session.mount("https://", agentrail_python.requests.adapter(agentrail))``.
+Each redirect is sent through the adapter again, so every actual destination is
+admitted on its own. The call's step ends when ``send`` returns: with
+``stream=True`` it does not wait for the body.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+from urllib.parse import urlsplit
+
+from requests import PreparedRequest, Response
+from requests.adapters import HTTPAdapter
+
+from ._handoff import OWNED_HEADERS
+from ._outbound import begin, finish, origin
+
+if TYPE_CHECKING:
+    from ._client import Client
+
+
+class AgentrailAdapter(HTTPAdapter):
+    def __init__(self, client: Client, **kwargs: Any) -> None:
+        super().__init__(**kwargs)
+        self._agentrail = client
+
+    def send(
+        self,
+        request: PreparedRequest,
+        stream: bool = False,
+        timeout: Any = None,
+        verify: bool | str = True,
+        cert: Any = None,
+        proxies: Any = None,
+    ) -> Response:
+        url = request.url or ""
+        parts = urlsplit(url)
+        target = request.path_url
+        outgoing = begin(
+            self._agentrail,
+            request.method or "GET",
+            origin(parts.scheme, parts.hostname or "", parts.port),
+            target.split("?", 1)[0],
+            target,
+            url,
+        )
+        if outgoing.strip:
+            for name in OWNED_HEADERS:
+                request.headers.pop(name, None)
+        request.headers.update(outgoing.headers)
+        if outgoing.span is None:
+            return super().send(request, stream, timeout, verify, cert, proxies)
+        try:
+            response = super().send(request, stream, timeout, verify, cert, proxies)
+        except BaseException as error:
+            finish(outgoing.span, None, error)
+            raise
+        finish(outgoing.span, response.status_code, None)
+        return response
+
+
+def adapter(client: Client, **kwargs: Any) -> AgentrailAdapter:
+    """An ``HTTPAdapter`` to mount on a ``requests.Session``; ``kwargs`` go to ``HTTPAdapter``."""
+    return AgentrailAdapter(client, **kwargs)

agentrail_python/wsgi.py

@@ -0,0 +1,158 @@
+# Copyright 2026 Agentrail. Licensed under the Apache License, Version 2.0.
+"""WSGI middleware: Flask, Django (WSGI), Falcon, Pyramid, Bottle."""
+
+from __future__ import annotations
+
+import contextvars
+import threading
+from collections.abc import Callable, Iterable, Iterator, Mapping
+from typing import TYPE_CHECKING, Any
+from urllib.parse import quote
+
+from opentelemetry import context as otel_context
+from opentelemetry.trace import Span, StatusCode
+
+from ._options import IncomingRequest
+
+if TYPE_CHECKING:
+    from ._client import Client
+
+StartResponse = Callable[..., Any]
+WSGIApp = Callable[[dict[str, Any], StartResponse], Iterable[bytes]]
+
+
+def _request(environ: Mapping[str, Any]) -> IncomingRequest:
+    headers: dict[str, str] = {}
+    for key, value in environ.items():
+        if key.startswith("HTTP_") and isinstance(value, str):
+            headers[key[5:].replace("_", "-").lower()] = value
+    for key in ("CONTENT_TYPE", "CONTENT_LENGTH"):
+        value = environ.get(key)
+        if isinstance(value, str) and value:
+            headers[key.replace("_", "-").lower()] = value
+    # Servers that keep the request line as sent (gunicorn, Werkzeug, uWSGI) give
+    # the exact target a handoff was signed over; PATH_INFO is already decoded.
+    raw = environ.get("RAW_URI") or environ.get("REQUEST_URI")
+    if isinstance(raw, str) and raw.startswith("/"):
+        target = raw
+    else:
+        path = quote(
+            str(environ.get("SCRIPT_NAME", "")) + str(environ.get("PATH_INFO", "/")),
+            safe="/%:@!$&'()*+,;=~",
+        )
+        query = str(environ.get("QUERY_STRING", ""))
+        target = f"{path}?{query}" if query else path
+    scheme = str(environ.get("wsgi.url_scheme", "http"))
+    host = headers.get("host") or f"{environ.get('SERVER_NAME', 'localhost')}:{environ.get('SERVER_PORT', '80')}"
+    address = environ.get("REMOTE_ADDR")
+    return IncomingRequest(
+        method=str(environ.get("REQUEST_METHOD", "GET")),
+        target=target,
+        scheme=scheme,
+        host=host,
+        headers=headers,
+        client=address if isinstance(address, str) and address else None,
+        raw=environ,
+    )
+
+
+class _Body:
+    """The response body; the entry span ends when the server closes it.
+
+    A streamed body runs after the application returned, in the context the
+    application ran in: its logs, milestones and outgoing calls belong to the
+    request, and a span it opens across chunks stays current between them. That
+    context is the body's own, so nothing it sets reaches the server's.
+    """
+
+    def __init__(
+        self,
+        body: Iterable[bytes],
+        finish: Callable[[BaseException | None], None],
+        run: contextvars.Context,
+    ) -> None:
+        self._body = body
+        self._finish = finish
+        self._run = run
+
+    def __iter__(self) -> Iterator[bytes]:
+        try:
+            chunks = self._run.run(iter, self._body)
+        except BaseException as error:
+            self._finish(error)
+            raise
+        while True:
+            try:
+                chunk = self._run.run(next, chunks)
+            except StopIteration:
+                return
+            except BaseException as error:
+                self._finish(error)
+                raise
+            yield chunk
+
+    def close(self) -> None:
+        close = getattr(self._body, "close", None)
+        if callable(close):
+            try:
+                self._run.run(close)
+            except BaseException as error:
+                self._finish(error)
+                raise
+        self._finish(None)
+
+
+class AgentrailWSGIMiddleware:
+    """Records each HTTP request the wrapped application serves as a journey step.
+
+    ``app.wsgi_app = agentrail.wsgi(app.wsgi_app)`` in Flask, or
+    ``application = agentrail.wsgi(get_wsgi_application())`` in Django.
+    """
+
+    def __init__(self, app: WSGIApp, *, client: Client) -> None:
+        self.app = app
+        self._client = client
+
+    def __call__(self, environ: dict[str, Any], start_response: StartResponse) -> Iterable[bytes]:
+        entered = self._client._enter(_request(environ))  # pyright: ignore[reportPrivateUsage]
+        if entered is None:
+            return self.app(environ, start_response)
+        span, ctx = entered
+        status: list[int] = []
+        ended = threading.Lock()
+
+        def finish(error: BaseException | None) -> None:
+            if not ended.acquire(blocking=False):
+                return
+            _end(span, status, error)
+
+        def observed_start(status_line: str, headers: Any, exc_info: Any = None) -> Any:
+            code = status_line.split(" ", 1)[0]
+            if code.isdigit():
+                status[:] = [int(code)]
+            return start_response(status_line, headers, exc_info) if exc_info else start_response(status_line, headers)
+
+        # The application and its body run in one context of their own, holding
+        # the request: what they set carries from the call to every chunk.
+        token = otel_context.attach(ctx)
+        try:
+            run = contextvars.copy_context()
+        finally:
+            otel_context.detach(token)
+        try:
+            body = run.run(self.app, environ, observed_start)
+        except BaseException as error:
+            finish(error)
+            raise
+        return _Body(body, finish, run)
+
+
+def _end(span: Span, status: list[int], error: BaseException | None) -> None:
+    if error is not None:
+        span.set_status(StatusCode.ERROR)
+        span.set_attribute("error.type", type(error).__qualname__)
+    if status:
+        span.set_attribute("http.response.status_code", status[0])
+        if status[0] >= 500:
+            span.set_status(StatusCode.ERROR)
+    span.end()