computer-agent-py 0.1.0__py3-none-any.whl

computeragent/telemetry/middleware/guardrails.py

@@ -0,0 +1,127 @@
+"""``GuardrailFilter`` — generic transform/drop middleware.
+
+Capabilities:
+
+* truncate over-long string values in payload (``max_attribute_length``)
+* drop tool-related events for tool names outside an allowlist (glob-aware)
+* enforce a running per-session cost ceiling
+* invoke a user-supplied async content filter
+
+The cost ceiling is enforced *across events* by tracking per-session running
+cost in a small in-memory map. The map shrinks on each ``session_ended`` so it
+doesn't grow without bound.
+"""
+
+from __future__ import annotations
+
+import fnmatch
+import logging
+from collections.abc import Awaitable, Callable
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from ..event import TelemetryEvent
+
+logger = logging.getLogger("computeragent.telemetry")
+
+ContentFilter = Callable[["TelemetryEvent"], Awaitable["TelemetryEvent | None"]]
+
+
+class GuardrailFilter:
+    """Apply size, scope, and cost guardrails before events reach sinks."""
+
+    def __init__(
+        self,
+        *,
+        max_attribute_length: int | None = 8192,
+        tool_name_allowlist: set[str] | None = None,
+        cost_ceiling_usd: float | None = None,
+        content_filter: ContentFilter | None = None,
+    ) -> None:
+        self.max_attribute_length = max_attribute_length
+        self.tool_name_allowlist = tool_name_allowlist
+        self.cost_ceiling_usd = cost_ceiling_usd
+        self.content_filter = content_filter
+        # session_id → running cost. Cleared on session_ended.
+        self._cost_by_session: dict[str, float] = {}
+        # Sessions that have already tripped the ceiling — drop further events
+        # from them until session_ended.
+        self._tripped: set[str] = set()
+
+    async def __call__(self, event: TelemetryEvent) -> TelemetryEvent | None:
+        # Remember the original kind so cleanup logic at the end can tell a
+        # real session_ended apart from a synthetic one we minted here.
+        original_kind = event.kind
+        synthesized_session_ended = False
+
+        # 1. Drop entirely if this session has already tripped the budget.
+        if event.session_id in self._tripped and event.kind != "session_ended":
+            return None
+
+        # 2. Tool-name allowlist: only applies to tool_use / tool_result.
+        if self.tool_name_allowlist is not None and event.kind in {"tool_use", "tool_result"}:
+            name = event.payload.get("tool_name") or event.payload.get("tool_use_id") or ""
+            if not _matches_any(name, self.tool_name_allowlist):
+                return None
+
+        # 3. Cost ceiling. usage_snapshot carries the running total cost.
+        if self.cost_ceiling_usd is not None and event.kind == "usage_snapshot":
+            cost = float(event.payload.get("total_cost_usd", 0.0) or 0.0)
+            self._cost_by_session[event.session_id] = cost
+            if cost > self.cost_ceiling_usd:
+                self._tripped.add(event.session_id)
+                # Transform this event into a synthetic session_ended so a
+                # sink consumer still sees a terminal event, then forward it.
+                event.kind = "session_ended"
+                event.payload = {
+                    "is_error": True,
+                    "subtype": "budget_exceeded",
+                    "total_cost_usd": cost,
+                    "cost_ceiling_usd": self.cost_ceiling_usd,
+                    "result": "",
+                    "num_turns": 0,
+                    "duration_ms": 0.0,
+                    "usage": event.payload.get("usage", {}),
+                }
+                synthesized_session_ended = True
+
+        # 4. Truncate long strings in payload (cheap; modifies in place).
+        if self.max_attribute_length is not None:
+            truncated = _truncate(event.payload, self.max_attribute_length)
+            assert isinstance(truncated, dict)  # _truncate preserves shape for dict in → dict out
+            event.payload = truncated
+
+        # 5. User-supplied filter — last call so it sees the final payload.
+        if self.content_filter is not None:
+            try:
+                filtered: TelemetryEvent | None = await self.content_filter(event)
+                if filtered is None:
+                    return None
+                event = filtered
+            except Exception:  # noqa: BLE001
+                logger.warning("content_filter raised; dropping event", exc_info=True)
+                return None
+
+        # 6. session_ended bookkeeping — only on a REAL terminal event.
+        # A synthetic session_ended we minted from the trip must keep
+        # _tripped state set, so further events from that session are dropped.
+        if original_kind == "session_ended" and not synthesized_session_ended:
+            self._cost_by_session.pop(event.session_id, None)
+            self._tripped.discard(event.session_id)
+
+        return event
+
+
+def _matches_any(name: str, allowlist: set[str]) -> bool:
+    """Glob match — ``mcp__*`` and ``Read``-style literals both supported."""
+    return any(fnmatch.fnmatchcase(name, pat) for pat in allowlist)
+
+
+def _truncate(value: object, limit: int) -> object:
+    if isinstance(value, dict):
+        return {k: _truncate(v, limit) for k, v in value.items()}
+    if isinstance(value, list):
+        return [_truncate(v, limit) for v in value]
+    if isinstance(value, str) and len(value) > limit:
+        return value[:limit] + f"…(+{len(value) - limit} chars)"
+    return value

computeragent/telemetry/middleware/pii.py

@@ -0,0 +1,158 @@
+"""``PiiRedactor`` — pattern-based redaction over ``event.payload``.
+
+Built-in patterns cover the most common shapes we don't want shipping to a
+remote observability backend by default: email, US phone, SSN, generic
+credit-card-shaped digit runs, and AWS access key IDs / secret keys. The user
+appends their own patterns via ``extra_patterns=`` and steers redaction by
+path with ``path_allowlist=`` / ``path_denylist=``.
+
+Fail-closed: a bad pattern or an exception inside the substitution drops the
+event rather than forwarding possibly-unredacted content.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import logging
+import re
+from typing import TYPE_CHECKING, Any, Literal
+
+if TYPE_CHECKING:
+    from ..event import TelemetryEvent
+
+logger = logging.getLogger("computeragent.telemetry")
+
+
+# Conservative built-in patterns — high precision, low recall. Users add their
+# own (badge IDs, internal account numbers, …) via extra_patterns.
+BUILTIN_PATTERNS: tuple[re.Pattern[str], ...] = (
+    # email — RFC 5322 simplified
+    re.compile(r"[A-Za-z0-9._%+\-]+@[A-Za-z0-9.\-]+\.[A-Za-z]{2,}"),
+    # US phone with optional country code; matches "+1 555-123-4567",
+    # "(555) 123-4567", "555.123.4567", etc.
+    re.compile(r"(?:\+?1[\s\-.])?\(?\d{3}\)?[\s\-.]\d{3}[\s\-.]\d{4}"),
+    # US SSN
+    re.compile(r"\b\d{3}-\d{2}-\d{4}\b"),
+    # Credit card shaped — 13–19 digit runs with optional separators.
+    re.compile(r"\b(?:\d[ \-]?){13,19}\b"),
+    # AWS access key ID
+    re.compile(r"\bAKIA[0-9A-Z]{16}\b"),
+)
+
+
+Strategy = Literal["mask", "hash", "drop"]
+
+
+class PiiRedactor:
+    """Walks ``event.payload`` and redacts matches in string values.
+
+    The walk descends into nested dicts and lists; tuples/sets are converted
+    to lists during processing then put back. Non-string leaves are left
+    alone. ``raw`` is never touched — sinks that need verbatim data can still
+    grab it, but content-shipping sinks (OTel, AgentOS) only read ``payload``.
+    """
+
+    def __init__(
+        self,
+        *,
+        strategy: Strategy = "mask",
+        extra_patterns: list[re.Pattern[str] | str] | None = None,
+        path_allowlist: list[str] | None = None,
+        path_denylist: list[str] | None = None,
+        include_builtins: bool = True,
+    ) -> None:
+        self.strategy = strategy
+        self._patterns: list[re.Pattern[str]] = []
+        if include_builtins:
+            self._patterns.extend(BUILTIN_PATTERNS)
+        for p in extra_patterns or []:
+            self._patterns.append(p if isinstance(p, re.Pattern) else re.compile(p))
+        self._allowlist = [p.strip(".") for p in (path_allowlist or [])]
+        self._denylist = [p.strip(".") for p in (path_denylist or [])]
+
+    async def __call__(self, event: TelemetryEvent) -> TelemetryEvent | None:
+        try:
+            event.payload = self._walk(event.payload, path="payload")
+        except Exception:  # noqa: BLE001
+            logger.warning("PiiRedactor failed; dropping event (fail-closed)", exc_info=True)
+            return None
+        return event
+
+    # ── walk ----------------------------------------------------------------
+
+    def _walk(self, value: Any, *, path: str) -> Any:
+        if isinstance(value, dict):
+            return {k: self._walk(v, path=f"{path}.{k}") for k, v in value.items()}
+        if isinstance(value, list):
+            return [self._walk(v, path=f"{path}[{i}]") for i, v in enumerate(value)]
+        if isinstance(value, tuple):
+            return tuple(self._walk(v, path=f"{path}[{i}]") for i, v in enumerate(value))
+        if isinstance(value, str):
+            return self._redact_if_applicable(value, path=path)
+        return value
+
+    # ── policy --------------------------------------------------------------
+
+    def _redact_if_applicable(self, value: str, *, path: str) -> str:
+        bare_path = path.removeprefix("payload.").removeprefix("payload")
+        # Denylist wins outright — force redaction (or drop the field entirely
+        # when strategy=drop, but at this layer we operate on the string value
+        # so 'drop' on a string is equivalent to empty string).
+        if any(_path_matches(bare_path, p) for p in self._denylist):
+            return self._apply_strategy(value, force=True)
+        # Allowlist short-circuits redaction — useful for whitelisted fields
+        # like a known-safe URL or filename.
+        if self._allowlist and any(_path_matches(bare_path, p) for p in self._allowlist):
+            return value
+        return self._apply_strategy(value, force=False)
+
+    def _apply_strategy(self, value: str, *, force: bool) -> str:
+        if force and self.strategy == "drop":
+            return ""
+
+        if self.strategy == "drop":
+            # In drop mode, if any pattern matches anywhere, return empty.
+            for pat in self._patterns:
+                if pat.search(value):
+                    return ""
+            return value
+
+        def replace(match: re.Match[str]) -> str:
+            captured = match.group(0)
+            if self.strategy == "hash":
+                digest = hashlib.sha256(captured.encode("utf-8")).hexdigest()[:12]
+                return f"<redacted:{digest}>"
+            return "****"
+
+        out = value
+        for pat in self._patterns:
+            try:
+                out = pat.sub(replace, out)
+            except re.error:
+                logger.warning("regex sub failed for pattern %r; skipping", pat.pattern)
+        # If the path is force-redacted but no pattern matched, return a
+        # blanket replacement so the field's content never leaks.
+        if force and out == value:
+            if self.strategy == "hash":
+                digest = hashlib.sha256(value.encode("utf-8")).hexdigest()[:12]
+                return f"<redacted:{digest}>"
+            return "****"
+        return out
+
+
+def _path_matches(path: str, pattern: str) -> bool:
+    """Naive glob-ish path matcher: supports literal segments and ``*``.
+
+    Matches ``payload.foo.bar`` against ``foo.bar``, ``foo.*``, ``*.bar``.
+    Not a full glob — intentionally simple so the rule semantics are obvious.
+    """
+    if pattern == "*":
+        return True
+    path_parts = path.split(".")
+    pat_parts = pattern.split(".")
+    if len(pat_parts) > len(path_parts):
+        return False
+    for pat, seg in zip(pat_parts, path_parts, strict=False):
+        if pat != "*" and pat != seg:
+            return False
+    return True

computeragent/telemetry/pipeline.py

@@ -0,0 +1,160 @@
+"""The telemetry pipeline — middleware then fan-out to sinks.
+
+Design properties:
+
+* **Never breaks an agent run.** Every middleware and every sink is wrapped
+  with broad exception handling; failures are logged and absorbed.
+* **Fail-closed for PII.** A middleware that raises drops the event (rather
+  than forwarding a possibly-unredacted one). Switchable via
+  ``fail_open=True`` for debug builds.
+* **Non-blocking sink emission.** Sinks run as background tasks so a slow
+  exporter never stalls the message hot path. ``flush()`` waits for them at
+  teardown.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import logging
+import weakref
+from typing import TYPE_CHECKING, Protocol
+
+if TYPE_CHECKING:
+    from collections.abc import Awaitable
+
+    from .event import TelemetryEvent
+
+logger = logging.getLogger("computeragent.telemetry")
+
+
+class TelemetryMiddleware(Protocol):
+    """Transform or drop an event. Return None to drop."""
+
+    async def __call__(self, event: TelemetryEvent) -> TelemetryEvent | None: ...
+
+
+class TelemetrySink(Protocol):
+    """Write an event to an external system."""
+
+    async def emit(self, event: TelemetryEvent) -> None: ...
+
+    async def flush(self, timeout: float = 5.0) -> None: ...
+
+    async def aclose(self) -> None: ...
+
+
+class Pipeline:
+    """Owns the middleware chain and sink fan-out.
+
+    A single ``Pipeline`` is normally enough per process; ``configure()`` builds
+    the module-level default. Tests construct their own.
+    """
+
+    def __init__(
+        self,
+        *,
+        middleware: list[TelemetryMiddleware] | None = None,
+        sinks: list[TelemetrySink] | None = None,
+        fail_open: bool = False,
+    ) -> None:
+        self._middleware: list[TelemetryMiddleware] = list(middleware or [])
+        self._sinks: list[TelemetrySink] = list(sinks or [])
+        self._fail_open = fail_open
+        # Track scheduled sink tasks so flush() can await them.
+        self._pending: weakref.WeakSet[asyncio.Task[None]] = weakref.WeakSet()
+
+    # ── configuration -------------------------------------------------------
+
+    def add_middleware(self, mw: TelemetryMiddleware) -> None:
+        self._middleware.append(mw)
+
+    def add_sink(self, sink: TelemetrySink) -> None:
+        self._sinks.append(sink)
+
+    def replace_middleware(self, mws: list[TelemetryMiddleware]) -> None:
+        self._middleware = list(mws)
+
+    def replace_sinks(self, sinks: list[TelemetrySink]) -> None:
+        self._sinks = list(sinks)
+
+    @property
+    def middleware(self) -> list[TelemetryMiddleware]:
+        return list(self._middleware)
+
+    @property
+    def sinks(self) -> list[TelemetrySink]:
+        return list(self._sinks)
+
+    # ── runtime -------------------------------------------------------------
+
+    async def emit(self, event: TelemetryEvent) -> None:
+        """Run middleware, then schedule sink emissions concurrently.
+
+        Returns as soon as scheduling is done; sink writes happen in the
+        background. Use :meth:`flush` to wait for them.
+        """
+        for mw in self._middleware:
+            try:
+                event_or_none = await mw(event)
+            except Exception:  # noqa: BLE001 - middleware must never break the agent
+                logger.warning(
+                    "telemetry middleware raised; %s",
+                    "forwarding event (fail_open=True)" if self._fail_open else "dropping event",
+                    exc_info=True,
+                )
+                if not self._fail_open:
+                    return
+                event_or_none = event
+            if event_or_none is None:
+                return
+            event = event_or_none
+
+        for sink in self._sinks:
+            task = asyncio.create_task(self._safe_emit(sink, event))
+            self._pending.add(task)
+
+    async def _safe_emit(self, sink: TelemetrySink, event: TelemetryEvent) -> None:
+        try:
+            await sink.emit(event)
+        except Exception:  # noqa: BLE001 - sinks must never break the agent
+            logger.warning("telemetry sink raised on emit (%s)", type(sink).__name__, exc_info=True)
+
+    async def flush(self, timeout: float = 5.0) -> None:
+        """Wait for outstanding sink writes, then flush each sink."""
+        # Snapshot the pending tasks; new ones added during await are picked up
+        # by sinks' own flush().
+        tasks: list[asyncio.Task[None]] = [t for t in self._pending if not t.done()]
+        if tasks:
+            try:
+                await asyncio.wait_for(
+                    asyncio.gather(*tasks, return_exceptions=True), timeout=timeout
+                )
+            except TimeoutError:
+                logger.warning(
+                    "telemetry flush timed out after %.1fs; %d sink tasks still pending",
+                    timeout,
+                    len([t for t in tasks if not t.done()]),
+                )
+
+        for sink in self._sinks:
+            try:
+                await sink.flush(timeout=timeout)
+            except Exception:  # noqa: BLE001
+                logger.warning(
+                    "telemetry sink raised on flush (%s)", type(sink).__name__, exc_info=True
+                )
+
+    async def aclose(self) -> None:
+        await self.flush()
+        for sink in self._sinks:
+            try:
+                await sink.aclose()
+            except Exception:  # noqa: BLE001
+                logger.warning(
+                    "telemetry sink raised on aclose (%s)", type(sink).__name__, exc_info=True
+                )
+
+
+def _await(awaitable: Awaitable[None]) -> Awaitable[None]:
+    """Type narrower for the linter (kept for future use)."""
+    return awaitable

computeragent/telemetry/sinks/__init__.py

@@ -0,0 +1,28 @@
+"""Telemetry sinks — OTel collector and AgentOS Mongo registry.
+
+Each sink is conditional on its extra. Import the sink class only after
+installing the corresponding extra; otherwise the import raises ``ImportError``
+with a clean install hint.
+"""
+
+from __future__ import annotations
+
+__all__ = ["AgentRegistrySink", "MongoMessageSink", "OtelSink"]
+
+
+def __getattr__(name: str):  # type: ignore[no-untyped-def]
+    """Lazy import so ``from .sinks import OtelSink`` raises a clean ImportError
+    only when the optional dep is actually missing."""
+    if name == "OtelSink":
+        from .otel import OtelSink as _OtelSink
+
+        return _OtelSink
+    if name == "AgentRegistrySink":
+        from .agentos import AgentRegistrySink as _AgentRegistrySink
+
+        return _AgentRegistrySink
+    if name == "MongoMessageSink":
+        from .message_archive import MongoMessageSink as _MongoMessageSink
+
+        return _MongoMessageSink
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")