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

computeragent/policy/types.py

@@ -0,0 +1,121 @@
+"""Canonical policy types — engine-agnostic dataclasses + the ``PolicyEngine``
+Protocol that OPA and Cedar implementations satisfy.
+
+The same :class:`PolicyInput` shape is sent to both engines so users can switch
+between OPA and Cedar without changing the worker code.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import Any, Literal, Protocol
+
+
+@dataclass(frozen=True)
+class PolicyPrincipal:
+    """Who is requesting the action — user / session / agent identifier."""
+
+    id: str
+    groups: list[str] = field(default_factory=list)
+    attributes: dict[str, Any] = field(default_factory=dict)
+
+    def to_json(self) -> dict[str, Any]:
+        return {
+            "id": self.id,
+            "groups": list(self.groups),
+            "attributes": dict(self.attributes),
+        }
+
+
+@dataclass(frozen=True)
+class PolicyAction:
+    """What is being attempted. v0.2 covers tool_use; ``name`` is extensible."""
+
+    name: Literal["tool_use"]
+    tool_name: str
+    tool_input: dict[str, Any] = field(default_factory=dict)
+
+    def to_json(self) -> dict[str, Any]:
+        return {
+            "name": self.name,
+            "tool_name": self.tool_name,
+            "tool_input": dict(self.tool_input),
+        }
+
+
+@dataclass(frozen=True)
+class PolicyResource:
+    """What the action is being attempted *on* — typically the running agent."""
+
+    type: str = "agent"
+    agent_name: str = ""
+    model: str = ""
+    session_id: str = ""
+    attributes: dict[str, Any] = field(default_factory=dict)
+
+    def to_json(self) -> dict[str, Any]:
+        return {
+            "type": self.type,
+            "agent_name": self.agent_name,
+            "model": self.model,
+            "session_id": self.session_id,
+            "attributes": dict(self.attributes),
+        }
+
+
+@dataclass(frozen=True)
+class PolicyInput:
+    """One authorization request. Identical shape across OPA and Cedar engines."""
+
+    principal: PolicyPrincipal
+    action: PolicyAction
+    resource: PolicyResource
+    context: dict[str, Any] = field(default_factory=dict)
+
+    def to_json(self) -> dict[str, Any]:
+        return {
+            "principal": self.principal.to_json(),
+            "action": self.action.to_json(),
+            "resource": self.resource.to_json(),
+            "context": dict(self.context),
+        }
+
+
+class PolicyDecision(Enum):
+    ALLOW = "allow"
+    DENY = "deny"
+
+
+@dataclass(frozen=True)
+class PolicyResult:
+    """Engine output — decision plus human-readable reason + optional obligations.
+
+    Obligations are advisory in v0.2 (not enforced); kept on the result so
+    sinks / audit logs can surface them.
+    """
+
+    decision: PolicyDecision
+    reason: str = ""
+    obligations: list[dict[str, Any]] = field(default_factory=list)
+    engine: str = ""  # populated by the concrete engine ("opa" | "cedar")
+
+
+class PolicyEngine(Protocol):
+    """Pluggable backend for tool-use authorization.
+
+    Implementations: :class:`~computeragent.policy.opa.OpaPolicyEngine` (remote
+    HTTP) and :class:`~computeragent.policy.cedar.CedarPolicyEngine` (in-process
+    via ``cedarpy``).
+    """
+
+    name: str
+    """Short identifier used in telemetry events (``"opa"`` / ``"cedar"``)."""
+
+    async def evaluate(self, input: PolicyInput) -> PolicyResult: ...
+
+    async def aclose(self) -> None: ...
+
+
+FailMode = Literal["deny", "allow"]
+"""What to return when the engine itself raises (network error, parse error)."""

computeragent/telemetry/__init__.py

@@ -0,0 +1,24 @@
+"""Telemetry surface — pipeline, middleware, sinks, configuration."""
+
+from __future__ import annotations
+
+from .config import configure, get_config, get_pipeline, shutdown
+from .event import EventKind, TelemetryEvent, new_session_id
+from .middleware.guardrails import GuardrailFilter
+from .middleware.pii import PiiRedactor
+from .pipeline import Pipeline, TelemetryMiddleware, TelemetrySink
+
+__all__ = [
+    "EventKind",
+    "GuardrailFilter",
+    "PiiRedactor",
+    "Pipeline",
+    "TelemetryEvent",
+    "TelemetryMiddleware",
+    "TelemetrySink",
+    "configure",
+    "get_config",
+    "get_pipeline",
+    "new_session_id",
+    "shutdown",
+]

computeragent/telemetry/config.py

@@ -0,0 +1,176 @@
+"""Module-level pipeline singleton + ``configure()`` API.
+
+Users either call :func:`configure` explicitly, or rely on the env-driven
+defaults that :func:`get_pipeline` builds on first use. Both paths are safe
+to interleave (later ``configure()`` replaces the lazy default).
+"""
+
+from __future__ import annotations
+
+import logging
+import os
+from typing import Any, Literal
+
+from .pipeline import Pipeline, TelemetryMiddleware, TelemetrySink
+
+logger = logging.getLogger("computeragent.telemetry")
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Globals
+# ──────────────────────────────────────────────────────────────────────────
+
+_PIPELINE: Pipeline | None = None
+_CONFIG: dict[str, Any] = {}
+
+
+def get_pipeline() -> Pipeline:
+    """Return the active pipeline, lazy-building from env if not configured."""
+    global _PIPELINE
+    if _PIPELINE is None:
+        _PIPELINE = _build_default_pipeline()
+    return _PIPELINE
+
+
+def configure(
+    *,
+    middleware: list[TelemetryMiddleware] | None = None,
+    sinks: list[TelemetrySink] | None = None,
+    service_name: str | None = None,
+    exporter: Literal["console", "otlp-http"] | None = None,
+    endpoint: str | None = None,
+    headers: dict[str, str] | None = None,
+    sample_rate: float = 1.0,
+    capture_content: bool = False,
+    capture_content_mode: Literal["events", "attributes", "both"] = "events",
+    max_attribute_length: int = 8192,
+    fail_open: bool = False,
+) -> Pipeline:
+    """Programmatic configuration. Idempotent — replaces the active pipeline.
+
+    Auto-attaches the bundled :class:`~computeragent.telemetry.sinks.otel.OtelSink`
+    and :class:`~computeragent.telemetry.sinks.agentos.AgentRegistrySink` if their
+    extras are installed and not already in ``sinks``. Pass ``sinks=[]`` to
+    disable all auto-attachment.
+    """
+    global _PIPELINE, _CONFIG
+    _CONFIG = {
+        "service_name": service_name or os.environ.get("OTEL_SERVICE_NAME", "computeragent"),
+        "exporter": exporter,
+        "endpoint": endpoint or os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"),
+        "headers": headers
+        if headers is not None
+        else _parse_headers(os.environ.get("OTEL_EXPORTER_OTLP_HEADERS")),
+        "sample_rate": sample_rate,
+        "capture_content": capture_content,
+        "capture_content_mode": capture_content_mode,
+        "max_attribute_length": max_attribute_length,
+    }
+
+    effective_sinks: list[TelemetrySink] = list(sinks) if sinks is not None else _auto_sinks()
+    _PIPELINE = Pipeline(
+        middleware=list(middleware or []),
+        sinks=effective_sinks,
+        fail_open=fail_open,
+    )
+    return _PIPELINE
+
+
+async def shutdown(timeout: float = 5.0) -> None:
+    """Flush and close the active pipeline. Safe to call multiple times."""
+    global _PIPELINE
+    if _PIPELINE is None:
+        return
+    try:
+        await _PIPELINE.aclose()
+    finally:
+        _PIPELINE = None
+
+
+# ──────────────────────────────────────────────────────────────────────────
+# Internals
+# ──────────────────────────────────────────────────────────────────────────
+
+
+def get_config() -> dict[str, Any]:
+    """Expose the resolved config to sinks (read-only contract)."""
+    return dict(_CONFIG)
+
+
+def _build_default_pipeline() -> Pipeline:
+    """Construct a pipeline from env vars only. No middleware by default."""
+    global _CONFIG
+    _CONFIG = {
+        "service_name": os.environ.get("OTEL_SERVICE_NAME", "computeragent"),
+        "exporter": None,
+        "endpoint": os.environ.get("OTEL_EXPORTER_OTLP_ENDPOINT"),
+        "headers": _parse_headers(os.environ.get("OTEL_EXPORTER_OTLP_HEADERS")),
+        "sample_rate": 1.0,
+        "capture_content": _env_bool("COMPUTERAGENT_CAPTURE_CONTENT", default=False),
+        "capture_content_mode": os.environ.get("COMPUTERAGENT_CAPTURE_CONTENT_MODE", "events"),
+        "max_attribute_length": int(os.environ.get("COMPUTERAGENT_MAX_ATTRIBUTE_LENGTH", "8192")),
+    }
+    return Pipeline(middleware=[], sinks=_auto_sinks(), fail_open=False)
+
+
+def _auto_sinks() -> list[TelemetrySink]:
+    """Build the sink list implied by environment variables.
+
+    OTel sink: on when ``COMPUTERAGENT_OTEL!=disabled`` (default on) AND the
+    ``opentelemetry-*`` packages are importable. Falls back to console
+    exporter when no OTLP endpoint is configured.
+
+    AgentOS sink: on when ``AGENTOS_MONGO_URL`` is set AND ``motor`` is
+    importable.
+    """
+    sinks: list[TelemetrySink] = []
+
+    if os.environ.get("COMPUTERAGENT_OTEL", "").lower() != "disabled":
+        try:
+            from .sinks.otel import OtelSink  # noqa: WPS433
+
+            sinks.append(OtelSink())
+        except ImportError:
+            logger.debug("OtelSink not attached (computer-agent-py[otel] not installed)")
+
+    mongo_url = os.environ.get("AGENTOS_MONGO_URL")
+    if mongo_url:
+        try:
+            from .sinks.agentos import AgentRegistrySink  # noqa: WPS433
+
+            sinks.append(AgentRegistrySink(mongo_url=mongo_url))
+        except ImportError:
+            logger.debug(
+                "AgentRegistrySink not attached (computer-agent-py[agentos] not installed)"
+            )
+
+    return sinks
+
+
+def _parse_headers(raw: str | None) -> dict[str, str]:
+    if not raw:
+        return {}
+    out: dict[str, str] = {}
+    for pair in raw.split(","):
+        if "=" not in pair:
+            continue
+        k, _, v = pair.partition("=")
+        k = k.strip()
+        v = v.strip()
+        if k:
+            out[k] = v
+    return out
+
+
+def _env_bool(name: str, *, default: bool) -> bool:
+    raw = os.environ.get(name)
+    if raw is None:
+        return default
+    return raw.strip().lower() in {"1", "true", "yes", "on"}
+
+
+def _reset_for_tests() -> None:
+    """Test hook — drop the cached pipeline so the next call rebuilds it."""
+    global _PIPELINE, _CONFIG
+    _PIPELINE = None
+    _CONFIG = {}

computeragent/telemetry/event.py

@@ -0,0 +1,355 @@
+"""``TelemetryEvent`` — the unit that flows through the telemetry pipeline.
+
+One event is emitted for each meaningful thing that happens during a query:
+session start, each message yielded, each tool call, the final usage snapshot,
+and session end. Middleware can mutate ``payload`` (the user-visible projection)
+or drop the event entirely; sinks consume the event and write it somewhere.
+
+The ``raw`` reference points back to the original upstream object so sinks
+that need the verbatim shape (e.g. a debug audit log) can grab it. Sinks that
+ship content to external systems should read from ``payload`` so middleware
+redaction is honoured.
+"""
+
+from __future__ import annotations
+
+import logging
+import uuid
+from dataclasses import dataclass, field
+from datetime import datetime, timezone
+from typing import Any, Literal
+
+logger = logging.getLogger("computeragent.telemetry")
+
+EventKind = Literal[
+    "session_started",
+    "user_message",
+    "assistant_message",
+    "system_message",
+    "tool_use",
+    "tool_result",
+    "usage_snapshot",
+    "session_ended",
+    "policy_decision",
+]
+
+
+def _utcnow() -> datetime:
+    return datetime.now(timezone.utc)
+
+
+@dataclass
+class TelemetryEvent:
+    """Single unit of telemetry. Mutable: middleware modifies ``payload`` in place."""
+
+    kind: EventKind
+    session_id: str
+    timestamp: datetime = field(default_factory=_utcnow)
+    agent_name: str | None = None
+    payload: dict[str, Any] = field(default_factory=dict)
+    raw: Any = None
+
+    # ──────────────────────────────────────────────────────────────────────
+    # Factories — keep all the upstream-shape-to-event-shape mapping in one
+    # place so message_translate-style spaghetti doesn't end up scattered.
+    # ──────────────────────────────────────────────────────────────────────
+
+    @classmethod
+    def session_started(
+        cls,
+        *,
+        session_id: str,
+        agent_name: str | None,
+        options: Any,
+        prompt: Any,
+    ) -> TelemetryEvent:
+        # Capture only metadata that a sink might want — never the full options
+        # object, which can hold callables (hooks, can_use_tool) that don't
+        # serialize.
+        # `system_prompt` may be a string OR a preset/file object — only persist
+        # the string form, since the inline-source manifest we build downstream
+        # treats it as plain CLAUDE.md content.
+        sp = getattr(options, "system_prompt", None)
+        system_prompt = sp if isinstance(sp, str) else None
+        payload: dict[str, Any] = {
+            "model": getattr(options, "model", None),
+            "system_prompt": system_prompt,
+            "permission_mode": getattr(options, "permission_mode", None),
+            "allowed_tools": list(getattr(options, "allowed_tools", []) or []),
+            "disallowed_tools": list(getattr(options, "disallowed_tools", []) or []),
+            "max_turns": getattr(options, "max_turns", None),
+            "cwd": str(getattr(options, "cwd", "") or ""),
+            "add_dirs": [str(p) for p in (getattr(options, "add_dirs", []) or [])],
+            "prompt": prompt if isinstance(prompt, str) else "<streaming>",
+        }
+        return cls(
+            kind="session_started",
+            session_id=session_id,
+            agent_name=agent_name,
+            payload=payload,
+            raw=options,
+        )
+
+    @classmethod
+    def session_ended_ok(
+        cls,
+        *,
+        session_id: str,
+        agent_name: str | None,
+        result: Any,
+        duration_ms: float,
+        last_assistant_text: str,
+    ) -> TelemetryEvent:
+        payload: dict[str, Any] = {
+            "is_error": bool(getattr(result, "is_error", False)),
+            "subtype": getattr(result, "subtype", "success"),
+            "result": getattr(result, "result", "") or last_assistant_text,
+            "num_turns": int(getattr(result, "num_turns", 0) or 0),
+            "total_cost_usd": float(getattr(result, "total_cost_usd", 0.0) or 0.0),
+            "duration_ms": duration_ms,
+            "duration_api_ms": float(getattr(result, "duration_api_ms", 0.0) or 0.0),
+            "usage": _extract_usage(result),
+        }
+        return cls(
+            kind="session_ended",
+            session_id=session_id,
+            agent_name=agent_name,
+            payload=payload,
+            raw=result,
+        )
+
+    @classmethod
+    def session_ended_error(
+        cls,
+        *,
+        session_id: str,
+        agent_name: str | None,
+        exc: BaseException,
+        duration_ms: float,
+    ) -> TelemetryEvent:
+        return cls(
+            kind="session_ended",
+            session_id=session_id,
+            agent_name=agent_name,
+            payload={
+                "is_error": True,
+                "subtype": "error",
+                "error_type": type(exc).__name__,
+                "error_message": str(exc),
+                "duration_ms": duration_ms,
+                "result": "",
+                "num_turns": 0,
+                "total_cost_usd": 0.0,
+                "usage": {},
+            },
+            raw=exc,
+        )
+
+    @classmethod
+    def policy_decision(
+        cls,
+        *,
+        input: Any,
+        result: Any,
+        error: BaseException | None,
+        latency_ms: float,
+        engine_name: str,
+    ) -> TelemetryEvent:
+        """Emitted by ``PolicyToolAuthorizer`` after each authorization check.
+
+        ``input`` is a :class:`computeragent.policy.PolicyInput` and ``result``
+        a :class:`PolicyResult` (or ``None`` if the engine raised). Sinks pick
+        this up via the standard pipeline; ``OtelSink`` annotates the active
+        ``execute_tool`` span with the decision and reason.
+        """
+        principal = getattr(input, "principal", None)
+        action = getattr(input, "action", None)
+        resource = getattr(input, "resource", None)
+        decision_value = getattr(getattr(result, "decision", None), "value", None)
+        payload: dict[str, Any] = {
+            "engine": engine_name,
+            "decision": decision_value or ("error" if error else "unknown"),
+            "reason": getattr(result, "reason", "") if result is not None else "",
+            "latency_ms": latency_ms,
+            "principal_id": getattr(principal, "id", ""),
+            "principal_groups": list(getattr(principal, "groups", []) or []),
+            "tool_name": getattr(action, "tool_name", ""),
+            "agent_name": getattr(resource, "agent_name", ""),
+            "session_id": getattr(resource, "session_id", "")
+            or getattr(input, "context", {}).get("session_id", ""),
+        }
+        if error is not None:
+            payload["error_type"] = type(error).__name__
+            payload["error_message"] = str(error)
+        return cls(
+            kind="policy_decision",
+            session_id=payload["session_id"] or "unknown",
+            agent_name=payload["agent_name"] or None,
+            payload=payload,
+            raw=result if result is not None else error,
+        )
+
+    @classmethod
+    def from_message(
+        cls, msg: Any, *, session_id: str, agent_name: str | None
+    ) -> list[TelemetryEvent]:
+        """Map one upstream message to zero or more telemetry events.
+
+        The kinds emitted depend on the upstream message type and any content
+        blocks it carries. The translation here is deliberately tolerant: if
+        upstream adds a new field we don't recognise, we just don't emit
+        anything for it rather than raising.
+        """
+        # Lazy import — only need the types at runtime when translation runs.
+        from claude_agent_sdk.types import (
+            AssistantMessage,
+            ResultMessage,
+            SystemMessage,
+            ToolResultBlock,
+            ToolUseBlock,
+            UserMessage,
+        )
+
+        events: list[TelemetryEvent] = []
+
+        if isinstance(msg, AssistantMessage):
+            text_parts: list[str] = []
+            for block in getattr(msg, "content", []) or []:
+                if isinstance(block, ToolUseBlock):
+                    events.append(
+                        cls(
+                            kind="tool_use",
+                            session_id=session_id,
+                            agent_name=agent_name,
+                            payload={
+                                "tool_use_id": getattr(block, "id", ""),
+                                "tool_name": getattr(block, "name", ""),
+                                "tool_input": dict(getattr(block, "input", {}) or {}),
+                            },
+                            raw=block,
+                        )
+                    )
+                else:
+                    text = getattr(block, "text", None)
+                    if isinstance(text, str):
+                        text_parts.append(text)
+            events.append(
+                cls(
+                    kind="assistant_message",
+                    session_id=session_id,
+                    agent_name=agent_name,
+                    payload={
+                        "model": getattr(msg, "model", None),
+                        "text": "".join(text_parts),
+                        "block_count": len(getattr(msg, "content", []) or []),
+                    },
+                    raw=msg,
+                )
+            )
+        elif isinstance(msg, UserMessage):
+            tool_results: list[dict[str, Any]] = []
+            text_parts = []
+            for block in getattr(msg, "content", []) or []:
+                if isinstance(block, ToolResultBlock):
+                    tool_results.append(
+                        {
+                            "tool_use_id": getattr(block, "tool_use_id", ""),
+                            "is_error": bool(getattr(block, "is_error", False)),
+                            "content": getattr(block, "content", ""),
+                        }
+                    )
+                else:
+                    text = getattr(block, "text", None)
+                    if isinstance(text, str):
+                        text_parts.append(text)
+            if tool_results:
+                for tr in tool_results:
+                    events.append(
+                        cls(
+                            kind="tool_result",
+                            session_id=session_id,
+                            agent_name=agent_name,
+                            payload=tr,
+                            raw=msg,
+                        )
+                    )
+            if text_parts:
+                events.append(
+                    cls(
+                        kind="user_message",
+                        session_id=session_id,
+                        agent_name=agent_name,
+                        payload={"text": "".join(text_parts)},
+                        raw=msg,
+                    )
+                )
+        elif isinstance(msg, SystemMessage):
+            events.append(
+                cls(
+                    kind="system_message",
+                    session_id=session_id,
+                    agent_name=agent_name,
+                    payload={
+                        "subtype": getattr(msg, "subtype", ""),
+                        "data": dict(getattr(msg, "data", {}) or {}),
+                    },
+                    raw=msg,
+                )
+            )
+        elif isinstance(msg, ResultMessage):
+            # ResultMessage also drives a `usage_snapshot` so token + cost
+            # telemetry doesn't only appear at session_ended (sinks that
+            # record metrics per session want the histogram measurement now).
+            events.append(
+                cls(
+                    kind="usage_snapshot",
+                    session_id=session_id,
+                    agent_name=agent_name,
+                    payload={
+                        "model": _model_from_options_or_result(msg),
+                        "total_cost_usd": float(getattr(msg, "total_cost_usd", 0.0) or 0.0),
+                        "usage": _extract_usage(msg),
+                    },
+                    raw=msg,
+                )
+            )
+        # Anything else (StreamEvent, RateLimitEvent, TaskStartedMessage, …):
+        # skip silently. These can be added later if a sink needs them.
+        return events
+
+
+def _extract_usage(result_like: Any) -> dict[str, int | float | None]:
+    """Pull tokens + cache fields out of a ResultMessage's nested usage shape."""
+    usage = getattr(result_like, "usage", None)
+    if usage is None:
+        return {}
+    # Pydantic-style models, dataclasses, or dicts — handle all three.
+    if hasattr(usage, "model_dump"):
+        try:
+            return dict(usage.model_dump())
+        except Exception:  # noqa: BLE001 - defensive
+            logger.debug("usage.model_dump() failed", exc_info=True)
+    if hasattr(usage, "__dataclass_fields__"):
+        return {k: getattr(usage, k, None) for k in usage.__dataclass_fields__}
+    if isinstance(usage, dict):
+        return dict(usage)
+    # Last resort — pick up the conventional Anthropic fields if they exist.
+    return {
+        k: getattr(usage, k, None)
+        for k in (
+            "input_tokens",
+            "output_tokens",
+            "cache_read_input_tokens",
+            "cache_creation_input_tokens",
+        )
+        if hasattr(usage, k)
+    }
+
+
+def _model_from_options_or_result(result: Any) -> str | None:
+    return getattr(result, "model", None)
+
+
+def new_session_id() -> str:
+    """Generate a placeholder session id used before the upstream returns one."""
+    return f"ca_{uuid.uuid4().hex}"

computeragent/telemetry/middleware/__init__.py

@@ -0,0 +1,8 @@
+"""Telemetry middleware — PII redaction and guardrail filtering."""
+
+from __future__ import annotations
+
+from .guardrails import GuardrailFilter
+from .pii import PiiRedactor
+
+__all__ = ["GuardrailFilter", "PiiRedactor"]