Blame for distinct_tools/trace.py

"""A bounded, JSON-safe record of every tool call an LLM agent attempted.

**Why this exists.** The DSPy agent runs inside the Deno sandbox; the tools it
calls run outside it, host-side, behind the broker.  Everything interesting
about a tool call therefore happens across a process boundary the agent's own
trajectory cannot see.  Without an explicit record, a denied call, a timed-out
call and a call that never happened are indistinguishable in the run log.

Three rules shape the design.

1. **A denial is a fact about the run, not an absence of one.** Every attempt
   is recorded, including attempts the broker refused before reaching a
   handler.  A user who selected two tools and got an answer that silently
   skipped one deserves to see why.
2. **Everything here is worker-supplied or tool-supplied, so everything here
   is untrusted.** Tool ids, error messages and argument values all originate
   outside the server.  Every string is bounded at construction and every
   value is a JSON scalar, so a consumer can escape once at the render
   boundary and be done.  See ``CONTRIBUTING.md`` honesty rule 4.
3. **Bounded at construction, never at render.** A trace that is only trimmed
   when displayed is a memory leak with a nice appearance.
"""

from __future__ import annotations

import json
import threading
from collections.abc import Mapping, Sequence
from dataclasses import dataclass, field
from types import MappingProxyType
from typing import Any

#: Longest recorded rendering of a call's arguments.  Large enough to read a
#: realistic call, small enough that a runaway argument cannot dominate a
#: result payload.
MAX_ARGUMENTS_CHARACTERS = 2_048

#: Longest recorded rendering of a successful result.  Deliberately shorter
#: than ``JobPolicy.max_output_bytes``: the trace is for diagnosis, and the
#: full result already travelled to the model.
MAX_RESULT_CHARACTERS = 1_024

MAX_MESSAGE_CHARACTERS = 500
MAX_TOOL_ID_CHARACTERS = 128
MAX_VERSION_CHARACTERS = 32
MAX_CODE_CHARACTERS = 80

#: Hard cap on recorded calls.  A job's call quota is far lower, so reaching
#: this means something is wrong; the cap stops it becoming a memory problem.
MAX_RECORDED_CALLS = 64

#: Where a refusal was decided.  Kept explicit because "the tool did not run"
#: has several very different causes and the operator needs to tell them apart.
OUTCOME_OK = "ok"
OUTCOME_TOOL_ERROR = "tool_error"
OUTCOME_DENIED_RUN = "denied_run_allowlist"
OUTCOME_DENIED_OPERATOR = "denied_operator_policy"
OUTCOME_QUOTA = "quota_exceeded"
OUTCOME_TIMEOUT = "timeout"
OUTCOME_INVALID = "invalid_input"

_OUTCOMES = frozenset(
    {
        OUTCOME_OK,
        OUTCOME_TOOL_ERROR,
        OUTCOME_DENIED_RUN,
        OUTCOME_DENIED_OPERATOR,
        OUTCOME_QUOTA,
        OUTCOME_TIMEOUT,
        OUTCOME_INVALID,
    }
)

#: Broker error codes mapped onto trace outcomes.  Anything unmapped becomes
#: ``OUTCOME_TOOL_ERROR``, which is the conservative reading: an unrecognised
#: failure is still a failure.
_CODE_TO_OUTCOME = MappingProxyType(
    {
        "tool_not_allowed": OUTCOME_DENIED_RUN,
        "operator_not_approved": OUTCOME_DENIED_OPERATOR,
        "quota_exceeded": OUTCOME_QUOTA,
        "timeout": OUTCOME_TIMEOUT,
        "invalid_input": OUTCOME_INVALID,
        "invalid_tool_ref": OUTCOME_INVALID,
        "invalid_output": OUTCOME_TOOL_ERROR,
        "output_too_large": OUTCOME_TOOL_ERROR,
        "host_not_allowed": OUTCOME_DENIED_OPERATOR,
        "tool_error": OUTCOME_TOOL_ERROR,
    }
)


def _clip(value: object, limit: int) -> str:
    """Return at most ``limit`` characters of ``value`` as plain text.

    Control characters are stripped rather than escaped.  The consumer still
    has to escape for its own output format; removing C0 here means a tool
    cannot smuggle a newline or a terminal control sequence into a log line.
    """

    text = "" if value is None else str(value)
    cleaned = "".join(character for character in text if ord(character) >= 32 or character == " ")
    if len(cleaned) <= limit:
        return cleaned
    return cleaned[:limit] + "...[truncated]"


def _render_json(value: object, limit: int) -> str:
    try:
        encoded = json.dumps(
            value, ensure_ascii=False, allow_nan=False, separators=(",", ":"), default=str
        )
    except (TypeError, ValueError):
        encoded = "<unrenderable>"
    return _clip(encoded, limit)


def _argument_keys(arguments: object) -> tuple[str, ...]:
    if not isinstance(arguments, Mapping):
        return ()
    keys = sorted(_clip(key, 64) for key in list(arguments)[:64])
    return tuple(keys)


@dataclass(frozen=True)
class ToolCallRecord:
    """One attempted tool call, whatever became of it.

    ``outcome`` is the field to branch on.  ``error_code`` preserves the
    broker's own vocabulary so a consumer that knows it loses nothing, but a
    consumer that does not can still tell a denial from a timeout.
    """

    sequence: int
    tool_id: str
    version: str
    outcome: str
    arguments_json: str = "{}"
    argument_keys: tuple[str, ...] = ()
    result_json: str = ""
    error_code: str | None = None
    error_message: str | None = None
    elapsed_ms: int = 0

    def __post_init__(self) -> None:
        if isinstance(self.sequence, bool) or not isinstance(self.sequence, int):
            raise ValueError("sequence must be an integer")
        if self.sequence < 1:
            raise ValueError("sequence must start at 1")
        if self.outcome not in _OUTCOMES:
            raise ValueError(f"unknown trace outcome {self.outcome!r}")
        object.__setattr__(self, "tool_id", _clip(self.tool_id, MAX_TOOL_ID_CHARACTERS))
        object.__setattr__(self, "version", _clip(self.version, MAX_VERSION_CHARACTERS))
        object.__setattr__(
            self, "arguments_json", _clip(self.arguments_json, MAX_ARGUMENTS_CHARACTERS)
        )
        object.__setattr__(self, "result_json", _clip(self.result_json, MAX_RESULT_CHARACTERS))
        object.__setattr__(
            self,
            "argument_keys",
            tuple(_clip(key, 64) for key in self.argument_keys[:64]),
        )
        if self.error_code is not None:
            object.__setattr__(self, "error_code", _clip(self.error_code, MAX_CODE_CHARACTERS))
        if self.error_message is not None:
            object.__setattr__(
                self, "error_message", _clip(self.error_message, MAX_MESSAGE_CHARACTERS)
            )
        if isinstance(self.elapsed_ms, bool) or not isinstance(self.elapsed_ms, int):
            raise ValueError("elapsed_ms must be an integer")
        if self.elapsed_ms < 0:
            raise ValueError("elapsed_ms must not be negative")

    @property
    def ref(self) -> str:
        return f"{self.tool_id}@{self.version}"

    @property
    def ok(self) -> bool:
        return self.outcome == OUTCOME_OK

    @property
    def denied(self) -> bool:
        return self.outcome in (OUTCOME_DENIED_RUN, OUTCOME_DENIED_OPERATOR)

    def to_dict(self) -> dict[str, Any]:
        """Return a flat, JSON-ready record with no nested containers.

        Flatness is deliberate.  A renderer that walks arbitrary nesting has
        to decide what to escape at every level; a flat record of bounded
        strings and integers has exactly one escaping rule.
        """

        value: dict[str, Any] = {
            "sequence": self.sequence,
            "tool_id": self.tool_id,
            "version": self.version,
            "ref": self.ref,
            "outcome": self.outcome,
            "ok": self.ok,
            "arguments_json": self.arguments_json,
            "argument_keys": list(self.argument_keys),
            "elapsed_ms": self.elapsed_ms,
        }
        if self.result_json:
            value["result_json"] = self.result_json
        if self.error_code is not None:
            value["error_code"] = self.error_code
        if self.error_message is not None:
            value["error_message"] = self.error_message
        return value

    def to_event(self) -> dict[str, Any]:
        """Return the narrower shape the existing ``tool_events`` field uses.

        Kept compatible on purpose: ``JobResult.tool_events`` already has
        consumers, and widening it is not this module's decision to make.
        """

        event: dict[str, Any] = {
            "tool_id": self.tool_id,
            "version": self.version,
            "ok": self.ok,
            "elapsed_ms": self.elapsed_ms,
        }
        if self.error_code is not None:
            event["error_code"] = self.error_code
        return event


@dataclass
class ToolTrace:
    """Ordered, bounded, thread-safe record of one job's tool calls.

    Thread safety matters: ``ToolBroker.call`` runs each handler on its own
    thread, and a DSPy agent may issue calls from more than one place.
    """

    job_id: str = ""
    _records: list[ToolCallRecord] = field(default_factory=list)
    _dropped: int = 0
    _lock: threading.Lock = field(default_factory=threading.Lock, repr=False)

    def __post_init__(self) -> None:
        self.job_id = _clip(self.job_id, MAX_TOOL_ID_CHARACTERS)

    def record(
        self,
        *,
        tool_id: str,
        version: str,
        outcome: str,
        arguments: object = None,
        result: object = None,
        error_code: str | None = None,
        error_message: str | None = None,
        elapsed_ms: int = 0,
    ) -> ToolCallRecord | None:
        """Append one attempt.  Returns ``None`` once the cap is reached."""

        with self._lock:
            if len(self._records) >= MAX_RECORDED_CALLS:
                self._dropped += 1
                return None
            sequence = len(self._records) + 1
            record = ToolCallRecord(
                sequence=sequence,
                tool_id=tool_id,
                version=version,
                outcome=outcome,
                arguments_json=_render_json(arguments if arguments is not None else {},
                                            MAX_ARGUMENTS_CHARACTERS),
                argument_keys=_argument_keys(arguments),
                result_json="" if result is None else _render_json(result, MAX_RESULT_CHARACTERS),
                error_code=error_code,
                error_message=error_message,
                elapsed_ms=max(0, int(elapsed_ms)) if not isinstance(elapsed_ms, bool) else 0,
            )
            self._records.append(record)
            return record

    def record_result(
        self,
        *,
        tool_id: str,
        version: str,
        arguments: object,
        result_envelope: Mapping[str, Any],
    ) -> ToolCallRecord | None:
        """Record from a ``ToolResult.to_dict()`` envelope.

        This is the path the broker uses, so the mapping from broker error
        codes to trace outcomes lives in exactly one place.
        """

        ok = result_envelope.get("ok") is True
        error = result_envelope.get("error")
        code = None
        message = None
        if isinstance(error, Mapping):
            code = error.get("code")
            message = error.get("message")
        outcome = OUTCOME_OK if ok else _CODE_TO_OUTCOME.get(str(code), OUTCOME_TOOL_ERROR)
        elapsed = result_envelope.get("elapsed_ms", 0)
        return self.record(
            tool_id=tool_id,
            version=version,
            outcome=outcome,
            arguments=arguments,
            result=result_envelope.get("output") if ok else None,
            error_code=None if code is None else str(code),
            error_message=None if message is None else str(message),
            elapsed_ms=elapsed if isinstance(elapsed, int) and not isinstance(elapsed, bool) else 0,
        )

    @property
    def records(self) -> tuple[ToolCallRecord, ...]:
        with self._lock:
            return tuple(self._records)

    @property
    def dropped(self) -> int:
        """Attempts not recorded because the cap was reached. Never hidden."""

        with self._lock:
            return self._dropped

    def denials(self) -> tuple[ToolCallRecord, ...]:
        return tuple(record for record in self.records if record.denied)

    def summary(self) -> dict[str, Any]:
        """Counts an operator can read at a glance, with no zero standing in
        for a missing measurement."""

        records = self.records
        by_outcome: dict[str, int] = {}
        for record in records:
            by_outcome[record.outcome] = by_outcome.get(record.outcome, 0) + 1
        return {
            "job_id": self.job_id,
            "attempted": len(records),
            "succeeded": sum(1 for record in records if record.ok),
            "denied": sum(1 for record in records if record.denied),
            "by_outcome": dict(sorted(by_outcome.items())),
            "dropped": self.dropped,
        }

    def to_dicts(self) -> tuple[dict[str, Any], ...]:
        return tuple(record.to_dict() for record in self.records)

    def to_events(self) -> tuple[dict[str, Any], ...]:
        """Return ``JobResult.tool_events``-shaped entries, denials included."""

        return tuple(record.to_event() for record in self.records)


def describe_for_agent(records: Sequence[ToolCallRecord]) -> str:
    """Render a trace back for the LLM agent to read.

    Used when a denial should change what the agent does next.  A model that
    is told plainly that a tool was refused can stop asking for it; a model
    given silence will try again and spend the budget.  The wording states the
    refusal as settled, because it is: the allowlist is immutable for the run.
    """

    if not records:
        return "No tools were called."
    lines = []
    for record in records:
        if record.ok:
            lines.append(f"{record.sequence}. {record.ref} succeeded.")
        elif record.outcome == OUTCOME_DENIED_RUN:
            lines.append(
                f"{record.sequence}. {record.ref} was refused: it is not in this run's "
                "allowed tools. This cannot change during the run; do not call it again."
            )
        elif record.outcome == OUTCOME_DENIED_OPERATOR:
            lines.append(
                f"{record.sequence}. {record.ref} was refused: this worker's operator has not "
                "approved it. This cannot change during the run; do not call it again."
            )
        elif record.outcome == OUTCOME_QUOTA:
            lines.append(
                f"{record.sequence}. {record.ref} was refused: the call budget is spent. "
                "Answer with what you already have."
            )
        else:
            code = record.error_code or record.outcome
            lines.append(f"{record.sequence}. {record.ref} failed ({code}).")
    return "\n".join(lines)


__all__ = [
    "MAX_ARGUMENTS_CHARACTERS",
    "MAX_RECORDED_CALLS",
    "MAX_RESULT_CHARACTERS",
    "OUTCOME_DENIED_OPERATOR",
    "OUTCOME_DENIED_RUN",
    "OUTCOME_INVALID",
    "OUTCOME_OK",
    "OUTCOME_QUOTA",
    "OUTCOME_TIMEOUT",
    "OUTCOME_TOOL_ERROR",
    "ToolCallRecord",
    "ToolTrace",
    "describe_for_agent",
]