toolwrit 0.2.0__py3-none-any.whl

toolwrit/toolwrit.py

@@ -0,0 +1,428 @@
+"""
+The runtime object an application holds.
+
+Toolwrit ties together the three pieces that have to agree with each other:
+the policy engine (may this happen), the ledger (is there capacity left) and
+the audit chain (what actually happened). Keeping them behind one object is
+what stops an integration from recording a decision it did not enforce.
+
+There are two guard methods rather than one. Python cannot await from a
+synchronous frame, so ``guard`` serves callers whose tool and approval handler
+are ordinary functions, and ``guard_async`` serves everyone else -- it awaits
+whatever it is given and accepts plain functions unchanged. Both run the same
+evaluate-record-enforce sequence; only the calling convention differs.
+"""
+
+from __future__ import annotations
+
+import asyncio
+import inspect
+import uuid
+from collections.abc import Mapping
+from typing import Any, Awaitable, Callable, Sequence, TypeVar
+
+from .audit.chain import AuditEntry, AuditLog, canonicalize
+from .budget.ledger import Ledger, TokenPrice, now_ms
+from ._js import js_locale_integer, js_number, js_round, js_to_fixed, to_jsonable
+from .policy.engine import evaluate
+from .types import (
+    BudgetUsage,
+    BudgetWarning,
+    Decision,
+    EvalContext,
+    Policy,
+    ToolCall,
+    Violation,
+)
+
+T = TypeVar("T")
+
+#: Called when a rule's effect is "ask". Return True to permit the call.
+#: Without a handler, "ask" is treated as a denial: an unattended run must not
+#: silently upgrade itself to allowed.
+ApprovalHandler = Callable[[ToolCall, Decision], "bool | Awaitable[bool]"]
+
+#: Warn thresholds used when a plan does not name its own.
+DEFAULT_WARN_AT = [0.8, 0.95]
+
+
+class ToolwritDenied(Exception):
+    """Thrown when a guarded call is refused. Catch this to feed the agent an error."""
+
+    def __init__(self, tool: str, decision: Decision, audit_hash: str) -> None:
+        super().__init__(f"toolwrit: {tool} denied — {decision.reason}")
+        self.tool = tool
+        self.decision = decision
+        #: Hash of the audit entry recording the refusal. Quote it in support tickets.
+        self.audit_hash = audit_hash
+
+
+def size_of(result: Any) -> int:
+    """
+    Size of a tool result, in bytes.
+
+    Deterministic by construction: the same result always measures the same, so
+    a replayed audit log reaches the same verdict. Structured results are
+    measured by their canonical JSON length, which equals what the TypeScript
+    implementation gets from JSON.stringify -- sorting keys reorders bytes but
+    does not add or remove any. A result that cannot be serialised falls back to
+    its string form and is therefore UNDER-counted: documented rather than
+    hidden, because under-counting a volume ceiling fails open.
+    """
+    if result is None:
+        return 0
+    if isinstance(result, str):
+        return len(result.encode("utf-8"))
+    if isinstance(result, (bytes, bytearray, memoryview)):
+        return len(bytes(result))
+
+    try:
+        return len(canonicalize(result).encode("utf-8"))
+    except Exception:
+        return len(str(result).encode("utf-8"))
+
+
+def _format(value: float) -> str:
+    """Readable numbers in warning text: 12500 -> "12,500", 1.5 -> "1.50"."""
+    return js_locale_integer(value) if float(value).is_integer() else js_to_fixed(value, 2)
+
+
+class Toolwrit:
+    def __init__(
+        self,
+        policy: Policy,
+        *,
+        run: str | None = None,
+        audit_file: str | None = None,
+        redact: Sequence[str] | None = None,
+        on_ask: ApprovalHandler | None = None,
+        on_warn: Callable[[BudgetWarning], None] | None = None,
+        now: Callable[[], int] | None = None,
+    ) -> None:
+        self._policy = policy
+        self._on_ask = on_ask
+        self._on_warn = on_warn
+        self.run = run or str(uuid.uuid4())
+        self._now = now or now_ms
+        self._ledger = Ledger()
+        self._audit = AuditLog(run=self.run, file=audit_file, redact=redact)
+        #: "dimension:threshold" keys already reported, so each fires exactly once.
+        self._warned: set[str] = set()
+        self._lock: asyncio.Lock | None = None
+        self._lock_loop: object | None = None
+
+        # The approved envelope goes into the chain before anything can consume it.
+        # Recording it as an ordinary entry keeps the chain uniform -- one shape,
+        # one verifier -- and means an operator cannot later claim a different
+        # budget was authorised than the one the run actually started under.
+        if policy.plan is not None:
+            plan = policy.plan
+            self._audit.record(
+                self._to_call(
+                    "toolwrit:plan",
+                    {
+                        "purpose": plan.purpose,
+                        "approvedBy": plan.approvedBy,
+                        "budget": to_jsonable(policy.budget)
+                        if policy.budget is not None
+                        else None,
+                        "warnAt": plan.warnAt
+                        if plan.warnAt is not None
+                        else list(DEFAULT_WARN_AT),
+                    },
+                ),
+                Decision(
+                    effect="allow",
+                    rule="plan",
+                    reason="run plan recorded",
+                    violations=[],
+                ),
+                {"calls": 0, "tokens": 0, "usd": 0, "bytes": 0},
+            )
+
+    def check(self, tool: str, args: dict[str, Any] | None = None) -> Decision:
+        """
+        Evaluate a call without recording or executing anything.
+        Use this for dry-runs and for showing an operator what a policy would do.
+        """
+        return evaluate(self._to_call(tool, args or {}), self._context())
+
+    def guard(
+        self,
+        tool: str,
+        args: dict[str, Any],
+        execute: Callable[[], T],
+    ) -> T:
+        """
+        Evaluate, record, and -- if permitted -- run ``execute``.
+
+        The decision is written to the audit chain before ``execute`` is invoked,
+        so a crash inside the tool still leaves evidence that the call was
+        authorised. A refusal raises ToolwritDenied rather than returning a
+        sentinel, because a silently-skipped side effect is the worst possible
+        failure mode here.
+        """
+        call = self._to_call(tool, args)
+        decision = evaluate(call, self._context())
+
+        if decision.effect == "ask":
+            approved = self._ask(call, decision)
+            if inspect.isawaitable(approved):
+                # Close it rather than leaking a never-awaited coroutine: the
+                # handler has not run any of its body yet, so nothing is lost.
+                getattr(approved, "close", lambda: None)()
+                raise TypeError(
+                    "the configured on_ask handler is asynchronous; use guard_async"
+                )
+            decision = self._resolve_ask(decision, approved)
+
+        self._enforce(tool, call, decision)
+        return self._measure(call, execute())
+
+    async def guard_async(
+        self,
+        tool: str,
+        args: dict[str, Any],
+        execute: Callable[[], T | Awaitable[T]],
+    ) -> T:
+        """``guard`` for async callers: awaits the approval handler and the tool."""
+        call = self._to_call(tool, args)
+
+        # Deciding, recording and counting is one critical section. Awaiting an
+        # approval handler used to leave it interruptible: concurrent calls all
+        # evaluated against the same pre-approval snapshot, so three guards
+        # against a budget of one were all approved and all executed. The tool
+        # itself still runs outside the lock, and the lock is only ever
+        # contended while an approval is genuinely pending.
+        async with self._decision_lock():
+            decision = evaluate(call, self._context())
+
+            if decision.effect == "ask":
+                approved = self._ask(call, decision)
+                if inspect.isawaitable(approved):
+                    approved = await approved
+                decision = self._resolve_ask(decision, approved)
+
+            self._enforce(tool, call, decision)
+
+            # A bytes ceiling is the one budget whose consumption is unknowable
+            # until after the tool has run, so concurrent calls would each be
+            # decided against a total none of them had contributed to yet. Where
+            # a ceiling is declared the tool runs inside the critical section
+            # too, bounding a run at the limit plus one call -- the bound the
+            # documentation promises. A policy with no bytes budget keeps full
+            # concurrency, so the cost falls only on the feature that needs it.
+            if self._policy.budget is not None and self._policy.budget.bytes is not None:
+                inside = execute()
+                if inspect.isawaitable(inside):
+                    inside = await inside
+                return self._measure(call, inside)  # type: ignore[arg-type]
+
+        result = execute()
+        if inspect.isawaitable(result):
+            result = await result
+        return self._measure(call, result)  # type: ignore[arg-type]
+
+    def meter(
+        self,
+        input_tokens: float,
+        output_tokens: float,
+        price: TokenPrice | Mapping[str, float] | None = None,
+    ) -> None:
+        """
+        Record model consumption against the budget. Call after every model turn.
+
+        ``price`` may be a TokenPrice or a plain mapping such as
+        ``{"input": 0.003, "output": 0.015}``.
+        """
+        self._ledger.add_usage(input_tokens, output_tokens, price, self._now())
+        self._report_progress()
+
+    def spend(self, usd: float) -> None:
+        """Record non-token spend, e.g. a metered third-party API call."""
+        self._ledger.add_spend(usd, self._now())
+        self._report_progress()
+
+    def usage(self) -> BudgetUsage:
+        """Current consumption. Useful for progress bars and for tests."""
+        return self._ledger.snapshot()
+
+    def entries(self) -> list[AuditEntry]:
+        """The audit chain recorded so far."""
+        return self._audit.all()
+
+    def head(self) -> str:
+        """Hash of the newest audit entry -- the receipt for this run so far."""
+        return self._audit.head()
+
+    # -- internals ---------------------------------------------------------- #
+
+    def _context(self) -> EvalContext:
+        return EvalContext(
+            policy=self._policy,
+            usage=self._ledger.snapshot(),
+            history=self._audit.history(),
+        )
+
+    def _enforce(self, tool: str, call: ToolCall, decision: Decision) -> None:
+        """Record the verdict, then act on it. Recording comes first, always."""
+        entry = self._audit.record(call, decision, self._usage_for_audit())
+        if decision.effect != "allow":
+            raise ToolwritDenied(tool, decision, entry.hash)
+        self._ledger.count_call(call.at)
+
+    def _decision_lock(self) -> asyncio.Lock:
+        """
+        The async critical section's lock, created on the running loop.
+
+        Built lazily because a Toolwrit may be constructed outside any event loop,
+        and an asyncio.Lock bound to the wrong loop is worse than none. The
+        synchronous ``guard`` does not take it: it has no await point, so it
+        cannot interleave with itself on one thread. Mixing ``guard`` and
+        ``guard_async`` on one instance is therefore not protected -- pick one.
+        """
+        loop = asyncio.get_running_loop()
+        if self._lock is None or self._lock_loop is not loop:
+            self._lock = asyncio.Lock()
+            self._lock_loop = loop
+        return self._lock
+
+    def _measure(self, call: ToolCall, result: T) -> T:
+        """
+        The volume half of containment: record how much the tool returned.
+
+        Attributed to the call's own timestamp rather than a fresh clock read,
+        so the ledger does not depend on how long the tool happened to take.
+        """
+        self._ledger.add_bytes(size_of(result), call.at)
+        self._report_progress()
+        return result
+
+    def _ask(self, call: ToolCall, decision: Decision) -> Any:
+        return None if self._on_ask is None else self._on_ask(call, decision)
+
+    def _resolve_ask(self, decision: Decision, approved: Any) -> Decision:
+        if self._on_ask is None:
+            return Decision(
+                effect="deny",
+                rule=decision.rule,
+                reason=f"{decision.reason} (no approval handler configured)",
+                violations=[
+                    Violation(
+                        rule=decision.rule if decision.rule is not None else "policy",
+                        constraint="ask",
+                        message=(
+                            "this call requires approval but no onAsk handler was "
+                            "configured; denying by default"
+                        ),
+                    )
+                ],
+            )
+
+        if approved:
+            return Decision(
+                effect="allow",
+                rule=decision.rule,
+                reason=f"{decision.reason} (approved)",
+                violations=list(decision.violations),
+            )
+        return Decision(
+            effect="deny",
+            rule=decision.rule,
+            reason=f"{decision.reason} (approval refused)",
+            violations=[
+                Violation(
+                    rule=decision.rule if decision.rule is not None else "policy",
+                    constraint="ask",
+                    message="approval was refused by the configured handler",
+                )
+            ],
+        )
+
+    def _report_progress(self) -> None:
+        """
+        Fire any warn thresholds the latest consumption has crossed.
+
+        Warnings are recorded in the audit chain as well as delivered to on_warn,
+        so "nobody told me it was at 95%" is answerable from the log rather than
+        from whether a Slack message happened to be delivered.
+        """
+        plan, budget = self._policy.plan, self._policy.budget
+        if plan is None or budget is None:
+            return
+
+        thresholds = plan.warnAt if plan.warnAt is not None else DEFAULT_WARN_AT
+        usage = self._ledger.snapshot()
+        elapsed = (
+            0 if usage.started_at is None else (self._now() - usage.started_at) / 1000
+        )
+
+        dimensions: list[tuple[str, float, float | None]] = [
+            ("calls", usage.calls, budget.calls),
+            ("tokens", usage.tokens, budget.tokens),
+            ("usd", usage.usd, budget.usd),
+            ("bytes", usage.bytes, budget.bytes),
+            ("seconds", elapsed, budget.seconds),
+        ]
+
+        for dimension, used, limit in dimensions:
+            if limit is None or limit <= 0:
+                continue
+            fraction = used / limit
+
+            for threshold in thresholds:
+                key = f"{dimension}:{js_number(threshold)}"
+                if fraction < threshold or key in self._warned:
+                    continue
+                self._warned.add(key)
+
+                warning = BudgetWarning(
+                    dimension=dimension,  # type: ignore[arg-type]
+                    threshold=threshold,
+                    used=used,
+                    limit=limit,
+                    message=(
+                        f'run "{self.run}" ({plan.purpose}) has used '
+                        f"{_format(used)}/{_format(limit)} {dimension} — "
+                        f"{js_round(fraction * 100)}% of the approved envelope"
+                    ),
+                )
+
+                self._audit.record(
+                    self._to_call(
+                        "toolwrit:warning",
+                        {
+                            "dimension": warning.dimension,
+                            "threshold": warning.threshold,
+                            "used": warning.used,
+                            "limit": warning.limit,
+                            "message": warning.message,
+                        },
+                    ),
+                    Decision(
+                        effect="allow",
+                        rule="plan",
+                        reason=warning.message,
+                        violations=[],
+                    ),
+                    {
+                        "calls": usage.calls,
+                        "tokens": usage.tokens,
+                        "usd": usage.usd,
+                        "bytes": usage.bytes,
+                    },
+                )
+                if self._on_warn is not None:
+                    self._on_warn(warning)
+
+    def _to_call(self, tool: str, args: dict[str, Any]) -> ToolCall:
+        return ToolCall(id=str(uuid.uuid4()), tool=tool, args=args, at=self._now())
+
+    def _usage_for_audit(self) -> dict[str, float]:
+        usage = self._ledger.snapshot()
+        return {
+            "calls": usage.calls,
+            "tokens": usage.tokens,
+            "usd": usage.usd,
+            "bytes": usage.bytes,
+        }

toolwrit/types.py

@@ -0,0 +1,262 @@
+"""
+Core contracts for Toolwrit.
+
+Everything in this file is data. The policy engine is a pure function over
+these shapes: no network, no model calls, no randomness. That is the whole
+point -- an enforcement decision must be reproducible from the policy, the
+call and the ledger state alone, so it can be replayed during an audit.
+
+Each dataclass declares ``JSON_OMIT_IF_NONE``: the optional fields that are
+simply *absent* from the JSON form rather than present-and-null. The
+distinction is load-bearing, because the audit chain hashes that JSON and
+``Decision.rule`` -- which is not in any omit set -- must serialise as
+``"rule":null`` to match the TypeScript chain.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, ClassVar, Literal, Sequence
+
+Effect = Literal["allow", "deny", "ask"]
+ArgType = Literal["string", "number", "boolean", "object", "array"]
+
+
+@dataclass
+class ToolCall:
+    """A single tool invocation an agent wants to make."""
+
+    #: Caller-supplied id, unique within a run. Used to correlate audit entries.
+    id: str
+    #: Tool name as the agent asked for it, e.g. "fs.write" or "github/create_issue".
+    tool: str
+    #: Arguments the agent proposed. Untrusted: this is model output.
+    args: dict[str, Any]
+    #: Milliseconds since epoch. Injected so evaluation stays deterministic in tests.
+    at: int
+
+
+@dataclass
+class BudgetLimits:
+    """What a run is allowed to consume before Toolwrit cuts it off."""
+
+    JSON_OMIT_IF_NONE: ClassVar[frozenset[str]] = frozenset(
+        {"calls", "tokens", "usd", "bytes", "seconds"}
+    )
+
+    #: Hard ceiling on total tool calls in the run.
+    calls: float | None = None
+    #: Hard ceiling on input+output tokens attributed to the run.
+    tokens: float | None = None
+    #: Hard ceiling on estimated spend, in USD.
+    usd: float | None = None
+    #: Hard ceiling on bytes returned BY tools across the run.
+    #: Bulk exfiltration is a permitted action repeated, not a forbidden one,
+    #: so only the running total gives it away. A result's size is unknowable
+    #: before the tool runs, so the breaching call completes and the NEXT one
+    #: is refused: this bounds a run at the limit plus one call.
+    bytes: float | None = None
+    #: Wall-clock ceiling for the run, in seconds, measured from the first call.
+    seconds: float | None = None
+
+
+@dataclass
+class BudgetUsage:
+    """Live consumption for a run, compared against BudgetLimits."""
+
+    calls: int = 0
+    tokens: float = 0
+    usd: float = 0
+    #: Bytes returned by tools so far.
+    bytes: float = 0
+    #: Milliseconds since epoch of the first metered event, or None before it.
+    started_at: int | None = None
+
+
+@dataclass
+class ArgConstraint:
+    """
+    Constraints applied to a single argument, addressed by a dotted path
+    (``path``, ``body.recipients.0``). Every constraint present must hold; an
+    absent value fails any constraint other than ``optional``.
+    """
+
+    JSON_OMIT_IF_NONE: ClassVar[frozenset[str]] = frozenset(
+        {
+            "optional", "oneOf", "noneOf", "matches", "startsWith", "excludes",
+            "min", "max", "maxLength", "minLength", "urlHosts", "type",
+        }
+    )
+
+    #: Allow the argument to be missing entirely. Default false.
+    optional: bool | None = None
+    #: Value must be one of these (deep-equality for objects).
+    oneOf: list[Any] | None = None
+    #: Value must NOT be one of these.
+    noneOf: list[Any] | None = None
+    #: String must match this regular expression (anchored by the author, not by us).
+    matches: str | None = None
+    #: String must start with one of these prefixes. Used for filesystem scoping.
+    startsWith: list[str] | None = None
+    #: String must NOT contain any of these substrings. Cheap traversal guard.
+    excludes: list[str] | None = None
+    #: Inclusive numeric bounds.
+    min: float | None = None
+    max: float | None = None
+    #: Length bounds for strings and arrays.
+    maxLength: float | None = None
+    minLength: float | None = None
+    #: Value must parse as a URL whose hostname is in this list (exact or ".suffix").
+    urlHosts: list[str] | None = None
+    #: Runtime type check.
+    type: ArgType | None = None
+
+
+@dataclass
+class RateLimit:
+    """Per-rule throttle, evaluated against the audit ledger."""
+
+    JSON_OMIT_IF_NONE: ClassVar[frozenset[str]] = frozenset({"perSeconds"})
+
+    #: Maximum matching calls allowed inside the window.
+    max: int = 0
+    #: Sliding window length in seconds. Omit for "per run".
+    perSeconds: float | None = None
+
+
+@dataclass
+class PolicyRule:
+    """One clause of a policy document."""
+
+    JSON_OMIT_IF_NONE: ClassVar[frozenset[str]] = frozenset(
+        {"description", "when", "limit"}
+    )
+
+    #: Stable identifier, surfaced in decisions and audit entries.
+    id: str
+    #: Tool name globs this rule applies to. ``*`` matches any run of characters
+    #: except ``.`` and ``/``; ``**`` matches anything, separators included.
+    #:
+    #: So ``**`` is the match-every-tool pattern. A bare ``*`` matches only names
+    #: with no separator in them -- it does NOT match "fs.read". Reach for ``**``
+    #: when you mean "everything", particularly on a deny rule, where a glob
+    #: that under-matches leaves a hole rather than failing closed.
+    tools: list[str] = field(default_factory=list)
+    effect: Effect = "allow"
+    #: Human-readable justification. Shows up in denial messages and reports.
+    description: str | None = None
+    #: Argument constraints, keyed by dotted path. All must hold for the rule to match.
+    when: dict[str, ArgConstraint] | None = None
+    #: Throttle for calls matching this rule. Exceeding it denies.
+    limit: RateLimit | None = None
+
+
+@dataclass
+class RunPlan:
+    """
+    The envelope a run declares before it starts, and a human approves once.
+
+    This is what makes capacity limits usable rather than annoying. A global
+    threshold ("50MB is suspicious") is either too loose to catch anything or
+    too tight to live with, because it has to guess at every job at once. A
+    declared envelope does not guess: a render job that asked for an hour and
+    two gigabytes gets exactly that without interruption, and the signal is not
+    "this looks like a lot" but "this run left the envelope its operator
+    approved". Approve at the start, hear nothing until a warn threshold.
+
+    The plan is written into the audit chain as its first entry, so what was
+    authorised is part of the tamper-evident record, not just what happened.
+    """
+
+    JSON_OMIT_IF_NONE: ClassVar[frozenset[str]] = frozenset({"approvedBy", "warnAt"})
+
+    #: What this run is for, in the operator's words. Carried into the audit log.
+    purpose: str
+    #: Who approved the envelope. Recorded, never verified by Toolwrit itself.
+    approvedBy: str | None = None
+    #: Fractions of the budget (0-1) at which the run reports to a human.
+    #: Each threshold fires at most once. Defaults to [0.8, 0.95].
+    warnAt: list[float] | None = None
+
+
+@dataclass
+class BudgetWarning:
+    """Emitted when consumption crosses one of the plan's warn thresholds."""
+
+    #: Which budget dimension crossed.
+    dimension: Literal["calls", "tokens", "usd", "bytes", "seconds"]
+    #: The threshold that fired, as a fraction of the limit.
+    threshold: float
+    #: Consumption and ceiling for that dimension.
+    used: float
+    limit: float
+    #: Ready-to-send summary, e.g. for a Slack message.
+    message: str
+
+
+@dataclass
+class Policy:
+    """A complete, self-contained enforcement policy."""
+
+    JSON_OMIT_IF_NONE: ClassVar[frozenset[str]] = frozenset(
+        {"name", "default", "budget", "plan"}
+    )
+
+    #: Schema version. Only "1" exists today; unknown versions are rejected.
+    version: Literal["1"] = "1"
+    rules: list[PolicyRule] = field(default_factory=list)
+    #: Free-text label carried into audit exports.
+    name: str | None = None
+    #: Effect when no rule matches. Defaults to "deny" -- Toolwrit is deny-by-default.
+    default: Effect | None = None
+    budget: BudgetLimits | None = None
+    #: The declared, pre-approved envelope for a run. See RunPlan.
+    plan: RunPlan | None = None
+
+
+@dataclass
+class Violation:
+    """Why a call was refused, or which constraint a rule turned on."""
+
+    JSON_OMIT_IF_NONE: ClassVar[frozenset[str]] = frozenset({"path"})
+
+    #: Rule that produced this violation, or "budget" / "policy" for engine-level ones.
+    rule: str
+    #: Constraint name that failed, e.g. "startsWith" or "usd".
+    constraint: str
+    #: Operator-facing explanation. Safe to show to the agent as a tool error.
+    message: str
+    #: Dotted argument path, when the violation is argument-specific.
+    path: str | None = None
+
+
+@dataclass
+class Decision:
+    """The engine's verdict on one tool call."""
+
+    effect: Effect
+    #: Id of the rule that decided this, or None when the default applied.
+    rule: str | None
+    #: Short reason, suitable for a log line.
+    reason: str
+    #: Empty for a clean allow.
+    violations: list[Violation] = field(default_factory=list)
+
+
+@dataclass(frozen=True)
+class HistoryEntry:
+    """A prior call, in the only shape the engine needs for rate limiting."""
+
+    tool: str
+    rule: str | None
+    at: int
+
+
+@dataclass
+class EvalContext:
+    """State the engine reads but never mutates."""
+
+    policy: Policy
+    usage: BudgetUsage = field(default_factory=BudgetUsage)
+    #: Prior calls in this run, oldest first. Used only for rate limiting.
+    history: Sequence[HistoryEntry] = field(default_factory=tuple)