"""Detach, bound and sanitize captured content before it is published.
This is the security boundary of the audit feature. Content is sanitized here,
*before* the action reaches the router queue -- the inspector is not the
boundary, and there is no "queue it raw and redact later" path.
The order is snapshot, then redact, then normalize again::
raw value -> bounded detached snapshot -> redactor -> bounded normalization
The application redactor therefore never receives a live business object, so
destructive code such as ``dict.pop()`` can only damage the observability copy.
The second pass exists because a redactor can return a new oversized, cyclic or
otherwise unsupported structure even when its input was safe.
What makes the snapshot an isolation boundary rather than a shallow copy is that
it accepts only *exact* built-in scalars and containers (plus ``UUID`` and
``datetime``, converted by the stdlib). An unknown object is a fixed sentinel: the normalizer never calls ``repr``, ``str``, ``__iter__``,
``keys``, ``asdict``, ``model_dump``, JSON hooks or any other
application-defined conversion on it, because those run arbitrary code -- on the
agent's hot path, on data the operator may have chosen not to capture at all.
"""
from __future__ import annotations
from datetime import datetime
from math import isfinite
from typing import Any, Final
from uuid import UUID
from edri.integrations.pydantic_ai.observability import (
NormalizationLimits,
Redactor,
)
__all__ = [
"BINARY",
"CYCLE",
"REDACTION_FAILED",
"TOO_DEEP",
"TRUNCATED",
"UNSUPPORTED",
"REQUEST_USAGE_FIELDS",
"USAGE_FIELDS",
"NormalizationBudget",
"capture_content",
"describe_error",
"normalize",
"usage_snapshot",
]
#: A value the normalizer will not expand: an application object, a container
#: subclass, a set, or a float that JSON cannot represent.
UNSUPPORTED: Final = "<unsupported>"
#: Raw binary. Fixed, with no length: a size is still information about content.
BINARY: Final = "<binary>"
#: The value is its own ancestor.
CYCLE: Final = "<cycle>"
#: Nesting passed ``max_depth``.
TOO_DEEP: Final = "<max-depth>"
#: A budget ran out -- value count, string length or total bytes.
TRUNCATED: Final = "<truncated>"
#: The application redactor raised. The original is *not* published instead.
REDACTION_FAILED: Final = "<redaction-failed>"
#: Approximate cost charged for a scalar, so a tree of millions of small values
#: still exhausts the byte budget rather than only the value count.
_SCALAR_BYTES: Final = 8
#: Exception attributes that may be published as structural metadata, and the
#: exact scalar types each is accepted as. Read by name from an allowlist: EDRI
#: never derives a code or a category by parsing free-form exception text.
_SAFE_ERROR_FIELDS: Final[tuple[tuple[str, tuple[type, ...]], ...]] = (
("code", (str, int)),
("status_code", (int,)),
("retryable", (bool,)),
)
#: Framework-owned classification, by exception type. Deliberately tiny: a
#: category is only reported when it follows from the type itself.
_ERROR_CATEGORIES: Final[tuple[tuple[type[BaseException], str], ...]] = (
(TimeoutError, "timeout"),
(ConnectionError, "connection"),
(MemoryError, "resource"),
)
#: Usage counters EDRI publishes, read explicitly so an upstream release adding
#: a field cannot silently change the audit shape. Measured against
#: ``pydantic_ai_slim`` 2.18.0; extending it is an intentional EDRI change.
USAGE_FIELDS: Final[tuple[str, ...]] = (
"requests",
"tool_calls",
"input_tokens",
"output_tokens",
"cache_write_tokens",
"cache_read_tokens",
"input_audio_tokens",
"cache_audio_read_tokens",
"output_audio_tokens",
)
#: Counters of *one model request*, which is a different measurement from the
#: run totals above: ``requests`` and ``tool_calls`` count whole-run activity
#: and a single request has no opinion about them. Listed separately rather
#: than filtered at the call site, so a future upstream release that adds a
#: run-scoped counter to the per-request object cannot leak it into a snapshot
#: that a consumer would then be tempted to sum with the run's own.
REQUEST_USAGE_FIELDS: Final[tuple[str, ...]] = (
"input_tokens",
"output_tokens",
"cache_write_tokens",
"cache_read_tokens",
"input_audio_tokens",
"cache_audio_read_tokens",
"output_audio_tokens",
)
[docs]
class NormalizationBudget:
"""Mutable spend against one action's limits.
One budget is shared by every field of one action -- request, response,
metadata, usage and error -- rather than each field getting the full
allowance, so an action cannot multiply the limit by its field count.
"""
__slots__ = ("limits", "values_remaining", "bytes_remaining")
def __init__(self, limits: NormalizationLimits | None = None) -> None:
self.limits = limits if limits is not None else NormalizationLimits()
self.values_remaining = self.limits.max_values
self.bytes_remaining = self.limits.max_total_bytes
[docs]
def spend(self, byte_count: int = _SCALAR_BYTES) -> bool:
"""Charge one value; ``False`` means it does not fit and nothing is spent."""
if self.values_remaining < 1 or self.bytes_remaining < byte_count:
return False
self.values_remaining -= 1
self.bytes_remaining -= byte_count
return True
[docs]
def fit_text(self, text: str) -> str | None:
"""Return as much of ``text`` as fits, or ``None`` if nothing does.
Truncation is by characters, never by bytes: slicing an encoded string
at a byte offset can split a codepoint, and the result has to stay valid
text. The byte budget is then applied by shrinking the character count
until the encoded form fits.
"""
if self.values_remaining < 1:
return None
was_empty = not text
text = text[: self.limits.max_string_characters]
encoded = len(text.encode("utf-8", "replace"))
while encoded > self.bytes_remaining and text:
# Assume the widest encoding for the overshoot, then re-measure;
# this converges in a couple of passes for any real payload.
keep = max(0, len(text) - max(1, (encoded - self.bytes_remaining) // 4))
text = text[:keep]
encoded = len(text.encode("utf-8", "replace"))
if not text and not was_empty:
# Nothing of it fits. An empty string would read as "the value was
# empty", which is a different claim from "it did not fit".
return None
self.values_remaining -= 1
self.bytes_remaining -= encoded
return text
#: Returned by :func:`_convert` when the budget cannot accept a value at all.
#: Distinct from a normalized ``TRUNCATED``: the caller stops the whole
#: container rather than storing a marker per remaining item.
_EXHAUSTED: Final = object()
class _Frame:
"""One container being filled by the iterative walk."""
__slots__ = ("items", "destination", "is_mapping", "depth", "ancestors")
def __init__(
self,
items: Any,
destination: Any,
is_mapping: bool,
depth: int,
ancestors: tuple[int, ...],
) -> None:
self.items = items
self.destination = destination
self.is_mapping = is_mapping
self.depth = depth
self.ancestors = ancestors
def store(self, key: Any, value: Any) -> None:
if self.is_mapping:
self.destination[key] = value
else:
self.destination.append(value)
def mark_truncated(self) -> None:
"""Record that the tail of this container was never captured.
One marker per container, not per dropped item: the point is that the
payload is incomplete, and repeating that per element would itself be
unbounded output.
"""
if self.is_mapping:
self.destination[TRUNCATED] = True
else:
self.destination.append(TRUNCATED)
def _convert(
value: Any, budget: NormalizationBudget, depth: int, ancestors: tuple[int, ...]
) -> tuple[Any, _Frame | None]:
"""Normalize one value, returning it plus a frame if it is a container."""
kind = type(value)
if kind is dict or kind is list or kind is tuple:
if id(value) in ancestors:
return (CYCLE, None) if budget.spend() else (_EXHAUSTED, None)
if depth >= budget.limits.max_depth:
return (TOO_DEEP, None) if budget.spend() else (_EXHAUSTED, None)
if not budget.spend(byte_count=2):
return _EXHAUSTED, None
nested = ancestors + (id(value),)
if kind is dict:
clone: Any = {}
# JSON objects are keyed by strings, and coercing a key would invent
# data the payload never held. A str *subclass* is an application
# type whose __str__ this module will not call either.
#
# Lazily, because the budget has to bound the work and not only the
# result: a tool returning a million-element payload under a
# two-value budget must cost two values, not a million-element copy
# of the payload taken before the first one is looked at.
items = (
(key, item) for key, item in value.items() if type(key) is str
)
return clone, _Frame(items, clone, True, depth, nested)
sequence: list[Any] = []
return sequence, _Frame(
((None, item) for item in value), sequence, False, depth, nested
)
if value is None or value is True or value is False:
return (value, None) if budget.spend() else (_EXHAUSTED, None)
if kind is str:
text = budget.fit_text(value)
return (text, None) if text is not None else (_EXHAUSTED, None)
if kind is int:
# An int has no width limit, so its decimal form is unbounded output
# even though it is a scalar.
if value.bit_length() > budget.limits.max_string_characters * 3:
return (TRUNCATED, None) if budget.spend() else (_EXHAUSTED, None)
return (value, None) if budget.spend() else (_EXHAUSTED, None)
if kind is float:
if not isfinite(value):
# NaN and the infinities are not JSON; a consumer that assumes they
# are would break on them.
return (UNSUPPORTED, None) if budget.spend() else (_EXHAUSTED, None)
return (value, None) if budget.spend() else (_EXHAUSTED, None)
if kind is bytes or kind is bytearray or kind is memoryview:
return (BINARY, None) if budget.spend() else (_EXHAUSTED, None)
if kind is UUID or kind is datetime:
# Identity, not payload: a run id, a parent action id, a dispatched
# event's ``_id``, a deadline. Degrading these to UNSUPPORTED would
# lose precisely the part of metadata the inspector correlates on.
# Exact types only, so the conversion is the stdlib's -- a subclass
# could override __str__/isoformat, and this module does not evaluate
# application conversions. Charged through fit_text like any other
# string, so the bound is the same one -- but all or nothing: half an
# identity still looks like an identity, and a consumer could correlate
# on the prefix and match the wrong record.
text = str(value) if kind is UUID else value.isoformat()
fitted = budget.fit_text(text)
if fitted != text:
return _EXHAUSTED, None
return fitted, None
return (UNSUPPORTED, None) if budget.spend() else (_EXHAUSTED, None)
[docs]
def normalize(value: Any, budget: NormalizationBudget) -> Any:
"""Return a detached, bounded, JSON-like clone of ``value``.
Total: no input raises, recurses without limit, or allocates without limit.
The emitter drops observations it cannot build, so a normalizer that raised
would silently delete audit records rather than produce a degraded one.
Iterative rather than recursive, because ``max_depth`` is caller-configurable
and a generous one would otherwise hit the interpreter's stack limit instead
of the documented bound.
"""
result, frame = _convert(value, budget, 0, ())
if result is _EXHAUSTED:
return TRUNCATED
stack = [frame] if frame is not None else []
while stack:
frame = stack[-1]
try:
key, item = next(frame.items)
except StopIteration:
stack.pop()
continue
except RuntimeError:
# The payload is still the application's own object while it is
# walked, so another thread can resize a container mid-traversal.
# The rest of it is not capturable, and this function does not
# raise: the snapshot says so and the walk moves on.
frame.mark_truncated()
stack.pop()
continue
child, child_frame = _convert(
item, budget, frame.depth + 1, frame.ancestors
)
if child is _EXHAUSTED:
frame.mark_truncated()
stack.pop()
continue
frame.store(key, child)
if child_frame is not None:
stack.append(child_frame)
return result
[docs]
def capture_content(
value: Any,
*,
budget: NormalizationBudget,
redactor: Redactor | None = None,
) -> Any:
"""Snapshot ``value``, hand the snapshot to ``redactor``, bound the result.
The redactor runs on the lifecycle hot path and is application code, so its
failure is contained here: the result degrades to
:data:`REDACTION_FAILED` and never falls back to the unredacted original,
which is the one outcome that would defeat the point of having a redactor.
"""
snapshot = normalize(value, budget)
if redactor is None:
return snapshot
try:
redacted = redactor(snapshot)
except Exception:
return REDACTION_FAILED
return normalize(redacted, budget)
[docs]
def usage_snapshot(
usage: Any,
budget: NormalizationBudget,
counters: tuple[str, ...] = USAGE_FIELDS,
) -> dict[str, Any] | None:
"""Build EDRI's stable usage shape from an explicit field list.
Counters that are absent are omitted rather than reported as zero: a missing
counter is a measurement EDRI did not make, and zero is a claim. Unknown
upstream fields are ignored so a dependency upgrade cannot change the
published schema.
``counters`` names the scope being snapshotted: :data:`USAGE_FIELDS` for a
whole run, :data:`REQUEST_USAGE_FIELDS` for one model request. The caller
picks it from what it is observing, because the two objects overlap and a
per-request object read through the run-scoped list would publish what looks
like a run total with counters missing.
Returns ``None`` when nothing supported is available -- "no usage" and "an
empty usage mapping" are different claims.
"""
if usage is None:
return None
snapshot: dict[str, Any] = {}
for name in counters:
try:
value = getattr(usage, name)
except Exception:
# A property can raise; that is not a reason to lose the rest.
continue
if type(value) is int:
snapshot[name] = value
try:
details = getattr(usage, "details", None)
except Exception:
details = None
if type(details) is dict and details:
normalized = normalize(details, budget)
if normalized:
snapshot["details"] = normalized
return snapshot or None
[docs]
def describe_error(
exception: BaseException,
*,
include_message: bool,
budget: NormalizationBudget,
redactor: Redactor | None = None,
) -> dict[str, Any]:
"""Describe ``exception`` under the effective content policy.
Exception text is not metadata-safe: provider messages carry request bodies,
model output, URLs with credentials and application rows, and truncation is
not sanitization. So ``metadata`` mode publishes only allowlisted structural
fields, and a message appears solely when the application enabled
``include_error_messages`` under ``content``.
The message is built from ``args``, never from ``str(exception)``: ``__str__``
is an application-defined conversion, and calling it to produce a nicer
message is exactly the arbitrary evaluation this module refuses. Neither the
live exception nor its traceback is retained -- a traceback holds frames,
locals and open resources, and an action outlives all of them.
"""
error: dict[str, Any] = {"type": type(exception).__qualname__}
for name, accepted in _SAFE_ERROR_FIELDS:
try:
value = getattr(exception, name, None)
except Exception:
continue
if value is None:
continue
# bool is a subclass of int, so an exact type check keeps a flag out of
# a numeric field and vice versa.
if type(value) in accepted:
error[name] = (
value[: budget.limits.max_string_characters]
if type(value) is str
else value
)
for kind, category in _ERROR_CATEGORIES:
if isinstance(exception, kind):
error["category"] = category
break
if include_message:
error["message"] = capture_content(
list(getattr(exception, "args", ())), budget=budget, redactor=redactor
)
return error