openscript-sdk 0.1.0__py3-none-any.whl

openscript_sdk/middleware/tool.py

@@ -0,0 +1,247 @@
+"""SecureTool — the policy pipeline around one tool, instead of around the agent.
+
+SecureAgent sees two moments: what the caller sent and what the agent finally
+returned. Everything an agent does in between — the tools it decides to call,
+with the arguments it decides to pass — happens inside ``agent.invoke()`` where
+no policy is looking. A prompt-injected agent does its damage there.
+
+SecureTool is the other half. It wraps a single tool callable and runs the same
+before / after pipeline on every call to it, with ``action="tool_call"``, so
+the checks that are about actions (ToolFirewallPolicy, approvals, audit) sit
+where the action is. It imports no agent framework: whatever calls the tool
+calls this instead.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import functools
+import inspect
+from collections.abc import Callable, Mapping, Sequence
+from typing import Any
+
+from openscript_sdk.contracts.interceptor import Policy
+from openscript_sdk.contracts.types import ActionContext
+from openscript_sdk.events.approvals import ApprovalStore
+from openscript_sdk.middleware.middleware import SecureAgent
+from openscript_sdk.observability.metrics import MetricsRecorder
+from openscript_sdk.observability.risk import RiskScorer
+from openscript_sdk.observability.tracing import ActionTracer
+
+_BY_POSITION_ONLY = (inspect.Parameter.POSITIONAL_ONLY, inspect.Parameter.VAR_POSITIONAL)
+
+
+class _ToolPipeline(SecureAgent):
+    """SecureAgent's phases with a tool in the middle instead of an agent.
+
+    A subclass rather than a second pipeline, so there is exactly one
+    definition of what a failure mode, an approval and a refusal mean.
+    """
+
+    async def run(
+        self, context: ActionContext, fn: Callable[..., Any], approval_id: str
+    ) -> ActionContext:
+        with self._trace(context):
+            context, granted = await self._run_before(context, approval_id)
+
+            # Read back from the context, not from what the caller passed: a
+            # policy may have rewritten an argument (a redacted credential),
+            # and the call that runs has to be the call that was checked.
+            args = context.input_data.get("args")
+            if not isinstance(args, dict):
+                raise RuntimeError(
+                    "a policy left the tool call without a mapping of arguments; refusing to run it"
+                )
+            if inspect.iscoroutinefunction(fn):
+                result = await fn(**args)
+            else:
+                # Off the event loop, as SecureAgent does for a sync agent. A
+                # callable object with an async __call__ comes back from the
+                # thread as a coroutine, which is awaited here.
+                result = await asyncio.to_thread(functools.partial(fn, **args))
+                if inspect.isawaitable(result):
+                    result = await result
+
+            context.output_data = {"output": result}
+            context = await self._run_after(context, approval_id, granted)
+            self._finalize(context)
+        return context
+
+
+class SecureTool:
+    """One tool, behind the same policies an agent gets.
+
+        send = SecureTool(send_email, policies=[ToolFirewallPolicy(rules_path="tools.yaml")])
+        await send(to="ops@example.com", body="done")     # called like the function
+
+    Each call becomes an ActionContext with ``action="tool_call"`` and
+    ``input_data={"name": ..., "args": {...}}``, where ``args`` holds every
+    argument by name *including the tool's own defaults* — a bound on ``amount``
+    is checked against the amount the tool will actually use, not skipped
+    because the caller left it out. (Which ones were defaults travels beside
+    the call, so a rule that refuses unlisted arguments refuses what a caller
+    passed and not what the tool's author wrote into its signature.)
+
+    A policy that denies in the before phase raises ActionBlockedError and the
+    tool does not run. One that requires approval raises it with an
+    ``approval_id``; once a human has approved, retry with
+    ``call(args, approval_id=...)``. The approval is single-use and bound to the
+    tool's name and exact arguments, so it cannot be spent on a different
+    recipient or a larger amount.
+
+    A deny in the after phase is a different thing and worth being clear about:
+    the tool has already run by then. What is withheld is its *result*, which
+    is how a tool's output is kept away from the model, not how its side
+    effects are undone.
+
+    A plain (non-async) tool runs in a worker thread so it cannot stall the
+    event loop.
+    """
+
+    def __init__(
+        self,
+        fn: Callable[..., Any],
+        policies: Sequence[Policy],
+        *,
+        name: str | None = None,
+        role: str | None = None,
+        agent_id: str = "default",
+        session_id: str = "default",
+        approval_store: ApprovalStore | None = None,
+        writer: Any | None = None,
+        risk_scorer: RiskScorer | None = None,
+        metrics: MetricsRecorder | None = None,
+        tracer: ActionTracer | None = None,
+    ) -> None:
+        if not callable(fn):
+            raise TypeError(f"SecureTool wraps a callable, got {type(fn).__name__}")
+        if not policies:
+            # SecureAgent falls back to a no-op here. A wrapper called
+            # SecureTool that lets every call through is not a default.
+            raise ValueError("SecureTool needs at least one policy")
+        resolved = name or getattr(fn, "__name__", "")
+        if not resolved:
+            raise ValueError("this tool has no __name__; pass name=")
+
+        try:
+            signature: inspect.Signature | None = inspect.signature(fn)
+        except (TypeError, ValueError):
+            # Builtins and some extension callables have none to read.
+            signature = None
+        if signature is not None:
+            for param in signature.parameters.values():
+                if param.kind in _BY_POSITION_ONLY:
+                    raise TypeError(
+                        f"{resolved}() takes '{param.name}' by position only. Rules are "
+                        "written against argument names, so every parameter of a guarded "
+                        "tool has to be passable by name."
+                    )
+
+        # Name, docstring and __wrapped__, so a framework that builds a tool
+        # schema by inspecting the callable sees the tool and not this class.
+        functools.update_wrapper(self, fn, updated=())
+        self.name = resolved
+        self._fn = fn
+        self._signature = signature
+        self._role = role
+        self._agent_id = agent_id
+        self._session_id = session_id
+        self._pipeline = _ToolPipeline(
+            agent=None,
+            policies=policies,
+            approval_store=approval_store,
+            writer=writer,
+            risk_scorer=risk_scorer,
+            metrics=metrics,
+            tracer=tracer,
+        )
+
+    async def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Call the tool the way the tool is called."""
+        result, _ = await self._run(*self._bind(args, kwargs), "", None, None, None)
+        return result
+
+    async def call(
+        self,
+        args: Mapping[str, Any] | None = None,
+        *,
+        approval_id: str = "",
+        role: str | None = None,
+        agent_id: str | None = None,
+        session_id: str | None = None,
+    ) -> Any:
+        """Call the tool with its arguments as a mapping.
+
+        The explicit form, for what a plain call has no room for: redeeming an
+        approval, or saying which role, agent and session the call belongs to
+        when one wrapped tool serves several.
+        """
+        result, _ = await self.call_with_context(
+            args, approval_id=approval_id, role=role, agent_id=agent_id, session_id=session_id
+        )
+        return result
+
+    async def call_with_context(
+        self,
+        args: Mapping[str, Any] | None = None,
+        *,
+        approval_id: str = "",
+        role: str | None = None,
+        agent_id: str | None = None,
+        session_id: str | None = None,
+    ) -> tuple[Any, ActionContext]:
+        """Like call(), but also returns the final ActionContext so callers can
+        read risk_score, risk_categories and per-policy metadata."""
+        return await self._run(*self._bind((), args or {}), approval_id, role, agent_id, session_id)
+
+    def _bind(
+        self, args: tuple[Any, ...], kwargs: Mapping[str, Any]
+    ) -> tuple[dict[str, Any], list[str]]:
+        """Every argument by name with the tool's defaults filled in, and the
+        names of the ones that came from a default rather than from the caller."""
+        if self._signature is None:
+            if args:
+                raise TypeError(
+                    f"{self.name}() has no readable signature, so its arguments "
+                    "have to be passed by name"
+                )
+            return dict(kwargs), []
+        # Raises the TypeError the tool itself would have, before any policy runs.
+        bound = self._signature.bind(*args, **kwargs)
+        passed = set(bound.arguments)
+        bound.apply_defaults()
+        named = dict(bound.arguments)
+        defaulted = [name for name in named if name not in passed]
+        for param in self._signature.parameters.values():
+            if param.kind is inspect.Parameter.VAR_KEYWORD:
+                # def tool(**options): the firewall should see the options, not
+                # one argument called "options".
+                named.update(named.pop(param.name))
+                defaulted = [name for name in defaulted if name != param.name]
+        return named, defaulted
+
+    async def _run(
+        self,
+        named: dict[str, Any],
+        defaulted: list[str],
+        approval_id: str,
+        role: str | None,
+        agent_id: str | None,
+        session_id: str | None,
+    ) -> tuple[Any, ActionContext]:
+        role = role if role is not None else self._role
+        # Beside the call, not in it: which arguments were defaulted is a fact
+        # about how the call was made, and must not change what an approval of
+        # the call is bound to.
+        metadata: dict[str, Any] = {"defaulted_args": defaulted}
+        if role is not None:
+            metadata["role"] = role
+        context = ActionContext(
+            action="tool_call",
+            agent_id=agent_id or self._agent_id,
+            session_id=session_id or self._session_id,
+            input_data={"name": self.name, "args": named},
+            metadata=metadata,
+        )
+        context = await self._pipeline.run(context, self._fn, approval_id)
+        return context.output_data.get("output"), context

openscript_sdk/observability/__init__.py

@@ -0,0 +1,11 @@
+from openscript_sdk.observability.metrics import MetricsRecorder
+from openscript_sdk.observability.risk import RiskScorer, aggregate_risk
+from openscript_sdk.observability.tracing import ActionTracer, tracer_from_env
+
+__all__ = [
+    "ActionTracer",
+    "MetricsRecorder",
+    "RiskScorer",
+    "aggregate_risk",
+    "tracer_from_env",
+]

openscript_sdk/observability/metrics.py

@@ -0,0 +1,105 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+if TYPE_CHECKING:
+    from openscript_sdk.contracts.types import ActionContext
+
+_RISK_BUCKETS = (0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0)
+
+
+class MetricsRecorder:
+    """Prometheus metrics for SecureAgent actions (pip install -e '.[metrics]').
+
+    Pass one instance per registry: SecureAgent(metrics=MetricsRecorder()).
+    Metric families:
+
+    - openscript_actions_total{decision}: every completed action, by final
+      decision (allow/deny/require_approval)
+    - openscript_action_risk_score: histogram of the aggregated risk_score
+    - openscript_policy_violations_total{policy}: one increment per policy
+      category that found something (metadata risk > 0) in an action
+    - openscript_injections_blocked_total: threat scorer flagged the input
+    - openscript_tool_calls_denied_total: tool firewall rejected a call
+    - openscript_pii_redacted_total: an action's output had PII redacted
+    - openscript_policy_latency_seconds{policy,phase}: per-hook latency
+      (recorded by SecureAgent around each before/after call)
+
+    Instantiating twice against the same registry raises (Prometheus forbids
+    duplicate timeseries) — share the instance, or pass a fresh
+    CollectorRegistry in tests.
+    """
+
+    def __init__(self, registry: Any | None = None) -> None:
+        try:
+            from prometheus_client import REGISTRY, Counter, Histogram
+        except ImportError as exc:  # pragma: no cover - exercised without the extra
+            raise ImportError(
+                "MetricsRecorder requires prometheus-client: pip install -e '.[metrics]'"
+            ) from exc
+
+        registry = registry if registry is not None else REGISTRY
+        self._actions = Counter(
+            "openscript_actions",
+            "Actions processed by SecureAgent, by final decision",
+            labelnames=("decision",),
+            registry=registry,
+        )
+        self._risk_score = Histogram(
+            "openscript_action_risk_score",
+            "Aggregated cross-policy risk score per action",
+            buckets=_RISK_BUCKETS,
+            registry=registry,
+        )
+        self._violations = Counter(
+            "openscript_policy_violations",
+            "Actions where a policy category found something (risk > 0)",
+            labelnames=("policy",),
+            registry=registry,
+        )
+        self._injections_blocked = Counter(
+            "openscript_injections_blocked",
+            "Actions flagged by the prompt-injection scorer",
+            registry=registry,
+        )
+        self._tool_calls_denied = Counter(
+            "openscript_tool_calls_denied",
+            "Tool calls rejected by the tool firewall",
+            registry=registry,
+        )
+        self._pii_redacted = Counter(
+            "openscript_pii_redacted",
+            "Actions whose output had PII redacted",
+            registry=registry,
+        )
+        self._policy_latency = Histogram(
+            "openscript_policy_latency_seconds",
+            "Latency of individual policy hooks",
+            labelnames=("policy", "phase"),
+            registry=registry,
+        )
+
+    def record_action(self, context: ActionContext) -> None:
+        """Record final counters for one action; called by SecureAgent at
+        every exit point, after risk aggregation."""
+        self._actions.labels(decision=context.decision.value).inc()
+        self._risk_score.observe(context.risk_score)
+
+        for category, risk in context.risk_categories.items():
+            if risk > 0:
+                self._violations.labels(policy=category).inc()
+
+        threat = context.metadata.get("threat")
+        if isinstance(threat, dict) and threat.get("flagged"):
+            self._injections_blocked.inc()
+
+        firewall = context.metadata.get("tool_firewall")
+        if isinstance(firewall, dict) and firewall.get("allowed") is False:
+            self._tool_calls_denied.inc()
+
+        pii = context.metadata.get("pii")
+        if isinstance(pii, dict) and pii.get("redacted"):
+            self._pii_redacted.inc()
+
+    def observe_policy_latency(self, policy: str, phase: str, seconds: float) -> None:
+        self._policy_latency.labels(policy=policy, phase=phase).observe(seconds)

openscript_sdk/observability/risk.py

@@ -0,0 +1,54 @@
+from __future__ import annotations
+
+from openscript_sdk.contracts.types import ActionContext
+
+
+def aggregate_risk(
+    context: ActionContext,
+    weights: dict[str, float] | None = None,
+) -> tuple[float, dict[str, float]]:
+    """Aggregate per-policy risk from standardized context.metadata entries.
+
+    Collects every metadata value shaped like {"risk": float, "category": str,
+    ...} (the shape all policies emit); anything else — plain kwargs, strings,
+    dicts without both keys — is ignored. When two entries share a category
+    the max wins. The total is the per-category sum, weighted by *weights*
+    (default 1.0 per category) and capped at 1.0.
+
+    Pure function of the context; safe to call at any point in the pipeline.
+    Returns (risk_score, risk_categories) where risk_categories maps
+    category -> unweighted per-category risk.
+    """
+    categories: dict[str, float] = {}
+    for entry in context.metadata.values():
+        if not isinstance(entry, dict):
+            continue
+        risk = entry.get("risk")
+        category = entry.get("category")
+        if not isinstance(risk, (int, float)) or isinstance(risk, bool):
+            continue
+        if not isinstance(category, str):
+            continue
+        clamped = min(1.0, max(0.0, float(risk)))
+        categories[category] = max(categories.get(category, 0.0), clamped)
+
+    weights = weights or {}
+    total = sum(risk * weights.get(category, 1.0) for category, risk in categories.items())
+    return round(min(1.0, total), 4), categories
+
+
+class RiskScorer:
+    """Writes the aggregated cross-policy risk back onto the ActionContext.
+
+    SecureAgent runs this after all policies (allow path) and before raising
+    ActionBlockedError (deny path), so every action carries a final
+    risk_score. *weights* scales individual categories' contribution to the
+    total (e.g. {"pii": 0.5} to down-weight redacted PII findings).
+    """
+
+    def __init__(self, weights: dict[str, float] | None = None) -> None:
+        self._weights = dict(weights or {})
+
+    def aggregate(self, context: ActionContext) -> ActionContext:
+        context.risk_score, context.risk_categories = aggregate_risk(context, self._weights)
+        return context

openscript_sdk/observability/tracing.py

@@ -0,0 +1,67 @@
+from __future__ import annotations
+
+import os
+from collections.abc import Iterator
+from contextlib import contextmanager
+from typing import TYPE_CHECKING, Any
+
+import structlog
+
+if TYPE_CHECKING:
+    from openscript_sdk.contracts.types import ActionContext
+
+logger = structlog.get_logger(__name__)
+
+# Feature flag: set OPENSCRIPT_OTEL=1 (or "true") to have SecureAgent open an
+# OpenTelemetry span per action without passing tracer= explicitly.
+OTEL_ENV_FLAG = "OPENSCRIPT_OTEL"
+
+
+class ActionTracer:
+    """One OpenTelemetry span per SecureAgent action (pip install -e '.[otel]').
+
+    Uses the ambient tracer provider — exporter setup (OTLP endpoint, which
+    Datadog and friends consume) is the host application's job via the
+    standard OTEL_* environment variables or explicit SDK configuration;
+    with no provider configured the spans are no-ops.
+    """
+
+    def __init__(self, instrumentation_name: str = "openscript") -> None:
+        try:
+            from opentelemetry import trace
+        except ImportError as exc:
+            raise ImportError(
+                "ActionTracer requires the OpenTelemetry SDK: pip install -e '.[otel]'"
+            ) from exc
+        self._tracer = trace.get_tracer(instrumentation_name)
+
+    @contextmanager
+    def span(self, context: ActionContext) -> Iterator[Any]:
+        with self._tracer.start_as_current_span(f"openscript.{context.action}") as span:
+            span.set_attribute("openscript.agent_id", context.agent_id)
+            span.set_attribute("openscript.session_id", context.session_id)
+            try:
+                yield span
+            finally:
+                # risk/decision are written by SecureAgent before any raise,
+                # so blocked actions carry their final values too
+                span.set_attribute("openscript.risk_score", context.risk_score)
+                span.set_attribute("openscript.decision", context.decision.value)
+
+
+def tracer_from_env() -> ActionTracer | None:
+    """Build an ActionTracer when the OPENSCRIPT_OTEL flag is set.
+
+    Fails open: with the flag set but the otel packages missing, logs a
+    warning and returns None rather than breaking the pipeline.
+    """
+    if os.environ.get(OTEL_ENV_FLAG, "").lower() not in ("1", "true"):
+        return None
+    try:
+        return ActionTracer()
+    except ImportError:
+        logger.warning(
+            "otel_flag_set_but_sdk_missing",
+            hint="pip install -e '.[otel]' to emit spans",
+        )
+        return None