marcus-toolkit 0.1.0__py3-none-any.whl

marcus_toolkit/usage.py

@@ -0,0 +1,222 @@
+"""
+Read what a run's agents spent, from the rows they wrote, while the run is going.
+
+Two layouts are read, because two kinds of agent exist today:
+
+* ``<state>/<agent>/<agent>/pi-sessions/**/*.jsonl``: the pi coding agent's
+  own transcripts, as the Minecraft harness writes them; every usage mapping
+  nested anywhere in an entry counts.
+* ``<state>/<agent>/usage.jsonl``: one row per model call, ``{"at": ..,
+  "model": .., "usage": {input, output, cacheRead, cacheWrite}}``, which any
+  runtime can write (ADR-0013 K8: a non-Claude-Code runtime has no transcript
+  to parse).
+
+A file may be mid-write when it is read during a live run: a trailing partial
+line is skipped, so this is safe to call repeatedly.
+"""
+
+from __future__ import annotations
+
+import json
+from collections.abc import Iterable, Iterator
+from pathlib import Path
+from typing import Any
+
+from . import pricing
+
+_ANTHROPIC_NAMES = {
+    "input_tokens": "input",
+    "output_tokens": "output",
+    "cache_read_input_tokens": "cacheRead",
+    "cache_creation_input_tokens": "cacheWrite",
+}
+
+
+def normalize_usage(raw: Any) -> dict[str, int] | None:
+    """
+    Read a usage mapping in either pi's or Anthropic's field names.
+
+    Parameters
+    ----------
+    raw : Any
+        A mapping with ``input``/``output`` or ``input_tokens``/``output_tokens``.
+
+    Returns
+    -------
+    Optional[Dict[str, int]]
+        ``{input, output, cacheRead, cacheWrite}`` with every key present, or
+        ``None`` if ``raw`` is not a usage mapping.
+    """
+    if not isinstance(raw, dict):
+        return None
+    mapped = {pricing.USAGE_KEYS[i]: 0 for i in range(4)}
+    seen = False
+    for key, value in raw.items():
+        name = _ANTHROPIC_NAMES.get(key, key)
+        if name in mapped and isinstance(value, int):
+            mapped[name] = value
+            if name in ("input", "output"):
+                seen = True
+    return mapped if seen else None
+
+
+def _usages(entry: Any) -> Iterable[dict[str, Any]]:
+    """Yield every usage mapping nested anywhere in a transcript entry."""
+    if isinstance(entry, dict):
+        if isinstance(entry.get("input"), int) and isinstance(entry.get("output"), int):
+            yield entry
+        for value in entry.values():
+            yield from _usages(value)
+    elif isinstance(entry, list):
+        for value in entry:
+            yield from _usages(value)
+
+
+def _lines(file: Path) -> Iterator[Any]:
+    for line in file.read_text(errors="replace").splitlines():
+        try:
+            yield json.loads(line)
+        except json.JSONDecodeError:
+            continue
+
+
+def rows(state_dir: Path) -> Iterator[tuple[str, dict[str, int], str | None]]:
+    """
+    Yield ``(agent, usage, model)`` for every usage row under ``state_dir``.
+
+    Parameters
+    ----------
+    state_dir : Path
+        The run's state directory, one subdirectory per agent.
+
+    Yields
+    ------
+    Tuple[str, Dict[str, int], Optional[str]]
+        The agent's name, its normalised usage, and the model the row names
+        (``None`` for pi transcripts, which do not say).
+    """
+    for session_dir in sorted(state_dir.glob("*/*/pi-sessions")):
+        agent = session_dir.parent.name
+        for file in sorted(session_dir.rglob("*.jsonl")):
+            for entry in _lines(file):
+                for usage in _usages(entry):
+                    normalized = normalize_usage(usage)
+                    if normalized is not None:
+                        yield agent, normalized, None
+    for file in sorted(state_dir.glob("*/usage.jsonl")):
+        agent = file.parent.name
+        for entry in _lines(file):
+            if not isinstance(entry, dict):
+                continue
+            normalized = normalize_usage(entry.get("usage"))
+            if normalized is not None:
+                model = entry.get("model")
+                yield agent, normalized, str(model) if model else None
+
+
+def costs(state_dir: Path, model_hint: str = "sonnet") -> dict[str, dict[str, Any]]:
+    """
+    Price each agent's usage.
+
+    Parameters
+    ----------
+    state_dir : Path
+        The run's state directory.
+    model_hint : str
+        The price list for rows that do not name their model.
+
+    Returns
+    -------
+    Dict[str, Dict[str, Any]]
+        Per agent: ``model_calls``, ``billed_calls`` (rows with any tokens;
+        pi records a zero row for an aborted continuation), token totals,
+        ``dollars`` and the models seen.
+    """
+    per_agent: dict[str, dict[str, Any]] = {}
+    for agent, usage, model in rows(state_dir):
+        row = per_agent.setdefault(
+            agent,
+            {
+                "model_calls": 0,
+                "billed_calls": 0,
+                "input": 0,
+                "output": 0,
+                "cache_read": 0,
+                "cache_write": 0,
+                "dollars": 0.0,
+                "models": [],
+            },
+        )
+        row["model_calls"] += 1
+        tokens = sum(usage.values())
+        if tokens > 0:
+            row["billed_calls"] += 1
+        row["input"] += usage["input"]
+        row["output"] += usage["output"]
+        row["cache_read"] += usage["cacheRead"]
+        row["cache_write"] += usage["cacheWrite"]
+        name = model or model_hint
+        row["dollars"] += pricing.price(usage, name)
+        key = pricing.resolve(name)
+        if key not in row["models"]:
+            row["models"].append(key)
+    for row in per_agent.values():
+        row["dollars"] = round(row["dollars"], 4)
+    return per_agent
+
+
+def total_dollars(state_dir: Path, model_hint: str = "sonnet") -> float:
+    """
+    Sum every agent's spend: the one number a stop-line check needs.
+
+    Parameters
+    ----------
+    state_dir : Path
+        The run's state directory.
+    model_hint : str
+        The price list for rows that do not name their model.
+
+    Returns
+    -------
+    float
+        Dollars so far, rounded to four places.
+    """
+    return float(
+        round(
+            sum(
+                float(agent["dollars"])
+                for agent in costs(state_dir, model_hint).values()
+            ),
+            4,
+        )
+    )
+
+
+def write_usage_row(
+    path: Path,
+    usage: dict[str, Any],
+    model: str,
+    at: str,
+    extra: dict[str, Any] | None = None,
+) -> None:
+    """
+    Append one usage row in the simple layout, for runtimes that write their own.
+
+    Parameters
+    ----------
+    path : Path
+        ``<state>/<agent>/usage.jsonl``; parents are created.
+    usage : Dict[str, Any]
+        Token counts in either naming.
+    model : str
+        The model billed.
+    at : str
+        An ISO timestamp.
+    extra : Optional[Dict[str, Any]]
+        Any other fields to keep beside the row (a turn number, a task id).
+    """
+    normalized = normalize_usage(usage) or {k: 0 for k in pricing.USAGE_KEYS}
+    path.parent.mkdir(parents=True, exist_ok=True)
+    row = {"at": at, "model": model, "usage": normalized, **(extra or {})}
+    with path.open("a") as handle:
+        handle.write(json.dumps(row) + "\n")

marcus_toolkit/views.py

@@ -0,0 +1,150 @@
+"""
+Who was shown a candidate, and who declined it.
+
+Evaporation counts views rather than seconds (Larry, 2026-11-06, "views,
+k=3"), and a view alone is not a decline: a decline requires an opportunity.
+An agent counts against a candidate only when, after the reply that showed
+it, that agent had at least one further turn in which support was possible
+and did not give it. The four cases:
+
+* shown, then supported: support, never a decline;
+* shown, then a turn without naming the subject: one decline;
+* shown, then no further turn (died, or the run ended): not a decline;
+* shown twice, then one turn: still one decline; agents are counted, not views.
+
+The data: the agent's runtime writes a ``candidate_view`` record into its
+``board-calls.jsonl`` for every reply that carried candidates, beside the one
+record per board call. Views are never reported to the board, where they
+could be mistaken for support.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from typing import Any
+
+from .candidates import ABSTAIN_MARK, OBJECT_MARK
+
+#: Calls only a deciding agent makes; an absorbed poll is not a turn.
+OPPORTUNITY_TOOLS: set[str] = {
+    "report_task_progress",
+    "give_opinion",
+    "log_artifact",
+    "get_task_context",
+    "query_project_history",
+    "log_decision",
+}
+#: Calls through which support can be given.
+SUPPORT_TOOLS: set[str] = {"report_task_progress", "give_opinion"}
+NOT_SUPPORT_MARKS = (OBJECT_MARK, ABSTAIN_MARK)
+
+
+def _is_support(what: Any) -> bool:
+    said = str(what or "").strip().upper()
+    return not any(said.startswith(mark) for mark in NOT_SUPPORT_MARKS)
+
+
+def tally(
+    calls: Iterable[dict[str, Any]], subjects: Iterable[str] | None = None
+) -> dict[str, dict[str, Any]]:
+    """
+    Compute, per candidate subject, who was shown it and who declined it.
+
+    Parameters
+    ----------
+    calls : Iterable[Dict[str, Any]]
+        Board-call records in time order (``marcus_toolkit.measure.read_calls``).
+    subjects : Optional[Iterable[str]]
+        Only these subjects; ``None`` tallies every subject seen.
+
+    Returns
+    -------
+    Dict[str, Dict[str, Any]]
+        ``{subject: {"shown_to", "declined_by", "supported_by", "awaiting_turn"}}``,
+        each a sorted list of agents.
+    """
+    wanted = set(subjects) if subjects is not None else None
+    first_view: dict[str, dict[str, str]] = {}
+    supported: dict[str, set[str]] = {}
+    turns: dict[str, list[str]] = {}
+
+    for call in calls:
+        agent = str(call.get("agent") or "")
+        tool = str(call.get("tool") or "")
+        at = str(call.get("at") or "")
+        args = call.get("args") or {}
+        if not agent or not at:
+            continue
+        if tool in OPPORTUNITY_TOOLS:
+            turns.setdefault(agent, []).append(at)
+        if tool == "candidate_view":
+            for subject in args.get("subjects") or []:
+                subject = str(subject)
+                if wanted is not None and subject not in wanted:
+                    continue
+                first_view.setdefault(subject, {}).setdefault(agent, at)
+        elif tool in SUPPORT_TOOLS:
+            subject = str(args.get("subject") or "")
+            if (
+                subject
+                and (wanted is None or subject in wanted)
+                and _is_support(args.get("what"))
+            ):
+                supported.setdefault(subject, set()).add(agent)
+
+    out: dict[str, dict[str, Any]] = {}
+    for subject in sorted(set(first_view) | set(wanted or ())):
+        seen = first_view.get(subject, {})
+        supporters = supported.get(subject, set())
+        declined: set[str] = set()
+        awaiting: set[str] = set()
+        for agent, view_at in seen.items():
+            if agent in supporters:
+                continue
+            if any(t > view_at for t in turns.get(agent, [])):
+                declined.add(agent)
+            else:
+                awaiting.add(agent)
+        out[subject] = {
+            "shown_to": sorted(seen),
+            "declined_by": sorted(declined),
+            "supported_by": sorted(supporters & set(seen)),
+            "awaiting_turn": sorted(awaiting),
+        }
+    return out
+
+
+def candidate_views(
+    calls: Iterable[dict[str, Any]], before: dict[str, str] | None = None
+) -> dict[str, list[str]]:
+    """
+    Count, per candidate subject, the agents that were shown it while it was open.
+
+    Parameters
+    ----------
+    calls : Iterable[Dict[str, Any]]
+        Board-call records.
+    before : Optional[Dict[str, str]]
+        Subject to its close time; a view at or after the close is a view of
+        a stale marker and does not count. A subject absent here is open.
+
+    Returns
+    -------
+    Dict[str, List[str]]
+        Subject to the sorted agents that saw it in time; subjects nobody saw
+        in time are absent.
+    """
+    seen: dict[str, set[str]] = {}
+    for call in calls:
+        if call.get("tool") != "candidate_view":
+            continue
+        agent = str(call.get("agent") or "")
+        at = str(call.get("at") or "")
+        for subject in (call.get("args") or {}).get("subjects") or []:
+            subject = str(subject)
+            closed = (before or {}).get(subject)
+            if closed and at >= closed:
+                continue
+            if agent:
+                seen.setdefault(subject, set()).add(agent)
+    return {s: sorted(a) for s, a in seen.items() if a}

marcus_toolkit/watchdog.py

@@ -0,0 +1,264 @@
+"""
+The cost watchdog: stop a run the moment its spend reaches the approved line.
+
+Lifted from the Minecraft plumbing's ``cost_watchdog.py`` (issue #815 there),
+with the defect the numerical verifier's run found fixed: the stop line is
+read fresh on every poll when it comes from a file, so a raised cap takes
+effect without a restart (``haiku-run-v2/REPORT.md``, section 8, item 4).
+
+The comparison is ``>=``: a stop line is the ceiling the human approved, and
+a run that has reached it has reached it. The watchdog returns on the poll
+that breaches and never polls again; the run it watched is expected to be
+gone.
+"""
+
+from __future__ import annotations
+
+import argparse
+import subprocess  # nosec B404 - stopping our own run's processes, fixed argv, no shell
+import sys
+import time
+from collections.abc import Callable, Sequence
+from dataclasses import dataclass
+from pathlib import Path
+
+from .usage import total_dollars
+
+StopLine = float | Callable[[], float]
+
+
+@dataclass
+class PollResult:
+    """What one poll found and did."""
+
+    dollars: float
+    stop_line: float
+    killed: bool
+
+
+class StopLineFile:
+    """
+    A stop line read from a file on every poll.
+
+    Parameters
+    ----------
+    path : Path
+        A file holding one number. A missing or unreadable file keeps the
+        last value read, and the first read must succeed.
+    """
+
+    def __init__(self, path: Path) -> None:
+        self.path = path
+        self._last = float(path.read_text().strip())
+
+    def __call__(self) -> float:
+        """Return the stop line, re-read from the file."""
+        try:
+            self._last = float(self.path.read_text().strip())
+        except (OSError, ValueError):
+            pass
+        return self._last
+
+
+def _line(stop_line: StopLine) -> float:
+    return float(stop_line()) if callable(stop_line) else float(stop_line)
+
+
+def watch_once(
+    state_dir: Path,
+    stop_line: StopLine,
+    kill: Callable[[], None],
+    model_hint: str = "sonnet",
+) -> PollResult:
+    """
+    Read accumulated cost once and kill the run if it has reached the stop line.
+
+    Parameters
+    ----------
+    state_dir : Path
+        The run's state directory.
+    stop_line : float or callable
+        The approved cap, or a callable that returns it fresh.
+    kill : Callable[[], None]
+        Called exactly once, synchronously, on breach.
+    model_hint : str
+        The price list for rows that do not name their model.
+
+    Returns
+    -------
+    PollResult
+    """
+    line = _line(stop_line)
+    dollars = total_dollars(state_dir, model_hint)
+    breached = dollars >= line
+    if breached:
+        kill()
+    return PollResult(dollars=dollars, stop_line=line, killed=breached)
+
+
+def watch(
+    state_dir: Path,
+    stop_line: StopLine,
+    kill: Callable[[], None],
+    poll_seconds: float = 30.0,
+    max_polls: int | None = None,
+    model_hint: str = "sonnet",
+    on_poll: Callable[[PollResult], None] | None = None,
+    sleep: Callable[[float], None] = time.sleep,
+) -> PollResult | None:
+    """
+    Poll until the run is killed or ``max_polls`` is exhausted.
+
+    Parameters
+    ----------
+    state_dir : Path
+        The run's state directory.
+    stop_line : float or callable
+        The approved cap, or a callable that returns it fresh each poll.
+    kill : Callable[[], None]
+        The stop action.
+    poll_seconds : float
+        Sleep between polls.
+    max_polls : Optional[int]
+        Stop after this many polls without a breach; ``None`` is unbounded.
+    model_hint : str
+        The price list for rows that do not name their model.
+    on_poll : Optional[Callable[[PollResult], None]]
+        Called after every poll, for logging.
+    sleep : Callable[[float], None]
+        The sleep function; tests pass a stub.
+
+    Returns
+    -------
+    Optional[PollResult]
+        The last poll, or ``None`` if none ran.
+    """
+    polls = 0
+    last: PollResult | None = None
+    while max_polls is None or polls < max_polls:
+        last = watch_once(state_dir, stop_line, kill, model_hint)
+        if on_poll:
+            on_poll(last)
+        polls += 1
+        if last.killed:
+            return last
+        if max_polls is not None and polls >= max_polls:
+            return last
+        sleep(poll_seconds)
+    return last
+
+
+def kill_action(
+    pid_files: Sequence[Path] = (), docker_name_pattern: str | None = None
+) -> Callable[[], None]:
+    """
+    Build the stop action: stop matching containers, then kill each pid in the files.
+
+    Each step is wrapped on its own, so one failing half does not block the
+    other: a half-stopped run still beats an unstopped one. Containers go
+    first, because they are what spends.
+
+    Parameters
+    ----------
+    pid_files : Sequence[Path]
+        Files holding one pid each (the server under test, the agents).
+    docker_name_pattern : Optional[str]
+        A ``docker ps --filter name=`` substring for the run's containers.
+
+    Returns
+    -------
+    Callable[[], None]
+    """
+
+    def _kill() -> None:
+        if docker_name_pattern:
+            try:
+                names = subprocess.run(  # nosec B603 B607
+                    [
+                        "docker",
+                        "ps",
+                        "--filter",
+                        f"name={docker_name_pattern}",
+                        "--format",
+                        "{{.Names}}",
+                    ],
+                    capture_output=True,
+                    text=True,
+                    timeout=10,
+                ).stdout.split()
+                for name in names:
+                    subprocess.run(
+                        ["docker", "stop", "-t", "5", name],
+                        capture_output=True,
+                        timeout=15,
+                    )  # nosec B603 B607
+            except Exception as error:  # noqa: BLE001
+                print(f"[watchdog] docker stop failed: {error}", file=sys.stderr)
+        for pid_file in pid_files:
+            try:
+                pid = int(Path(pid_file).read_text().strip())
+                subprocess.run(
+                    ["kill", str(pid)], capture_output=True, timeout=5
+                )  # nosec B603 B607
+            except Exception as error:  # noqa: BLE001
+                print(
+                    f"[watchdog] kill from {pid_file} failed: {error}", file=sys.stderr
+                )
+
+    return _kill
+
+
+def main(argv: list[str] | None = None) -> int:
+    """CLI entry point: the background watchdog a run script starts."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--state", required=True, help="the run's state directory")
+    group = parser.add_mutually_exclusive_group(required=True)
+    group.add_argument("--stop-line", type=float, help="approved hard dollar cap")
+    group.add_argument(
+        "--stop-line-file", help="a file holding the cap, re-read every poll"
+    )
+    parser.add_argument("--poll-seconds", type=float, default=30.0)
+    parser.add_argument(
+        "--model", default="sonnet", help="price list for rows that do not name a model"
+    )
+    parser.add_argument(
+        "--pid-file",
+        action="append",
+        default=[],
+        help="a file holding a pid to kill on breach (repeatable)",
+    )
+    parser.add_argument("--docker-name-pattern", default=None)
+    parser.add_argument("--max-minutes", type=float, default=None)
+    args = parser.parse_args(argv)
+
+    stop_line: StopLine = (
+        StopLineFile(Path(args.stop_line_file))
+        if args.stop_line_file
+        else float(args.stop_line)
+    )
+    max_polls = (
+        int(args.max_minutes * 60 / args.poll_seconds) if args.max_minutes else None
+    )
+
+    def log(result: PollResult) -> None:
+        tag = "BREACH -- stopping the run" if result.killed else "ok"
+        print(
+            f"[watchdog] ${result.dollars:.2f} of ${result.stop_line:.2f}"
+            f" stop line -- {tag}",
+            flush=True,
+        )
+
+    watch(
+        Path(args.state),
+        stop_line,
+        kill_action([Path(p) for p in args.pid_file], args.docker_name_pattern),
+        poll_seconds=args.poll_seconds,
+        max_polls=max_polls,
+        model_hint=args.model,
+        on_poll=log,
+    )
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())