@seanyao/roll 0.5.0 → 2.602.2

package/lib/agent_usage/pi.py

@@ -0,0 +1,200 @@
+"""
+pi agent usage extractor.
+
+pi runs in the loop as ``pi -p`` (text mode), whose stdout is ONLY the
+assistant's answer text — it carries no token/cost summary.  So stdout
+scraping (the ``extract()`` registry contract) cannot recover usage and
+always returns None for real pi output.
+
+Instead, pi persists every session to disk at::
+
+    ~/.pi/agent/sessions/<encoded-cwd>/<ISO-ts>_<uuid>.jsonl
+
+Each file is NDJSON: one ``{"type":"session","cwd":<abs-worktree-path>}``
+header line followed by ``{"type":"message","message":{...}}`` lines.
+Assistant messages carry a per-call ``usage`` block including pi's own
+cost calc.  The authoritative usage path is therefore ``usage_from_session``,
+which sums per-message usage for a cycle's worktree.  See ``pi_emit.py``
+(live capture) and ``backfill-pi-usage.py`` (historical backfill).
+"""
+
+import glob
+import json
+import os
+from typing import Optional
+
+
+def extract(stdin_lines: list[str]) -> Optional[dict]:
+    """Registry contract stub.
+
+    pi ``-p`` text-mode stdout carries no usage data, so this always
+    returns None and the caller falls back to the null-payload path.
+    Real usage is recovered from session files via ``usage_from_session``.
+    Kept so the agent_usage REGISTRY contract / tests stay valid.
+    """
+    return None
+
+
+# ── Session-file extraction (authoritative) ────────────────────────────────
+
+# pi reports a per-message ``cost.total``; we sum it into ``cost_reported``
+# for audit only.  The authoritative list cost is frozen by the writers from
+# lib/prices/snapshot-*-deepseek.json in deepseek's native currency (CNY) —
+# we never convert currencies (the CLI already shows the currency symbol).
+def _sessions_base_dir(base_dir: Optional[str]) -> str:
+    """Resolve the pi sessions root: arg → env → default."""
+    return (
+        base_dir
+        or os.environ.get("ROLL_PI_SESSIONS_DIR")
+        or os.path.expanduser("~/.pi/agent/sessions")
+    )
+
+
+def _sum_session_file(path: str) -> Optional[dict]:
+    """Sum per-message assistant usage in a single session jsonl.
+
+    Returns a usage dict (tokens summed) or None when the file has no
+    assistant usage.  Field mapping from pi → roll schema:
+    cacheWrite→cache_creation_tokens, cacheRead→cache_read_tokens.
+
+    ``cost_reported`` carries pi's own per-message ``cost.total`` summed,
+    purely for audit — it is NOT the authoritative cost.  The authoritative
+    list cost is frozen by the writers (pi_emit / backfill) from the deepseek
+    price snapshot in its native currency (CNY), matching claude's
+    ``_price_at_snapshot`` convention.  We never convert currencies.
+    """
+    tin = tout = tcr = tcw = 0
+    cost = 0.0
+    model = None
+    seen = False
+    try:
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    o = json.loads(line)
+                except json.JSONDecodeError:
+                    continue
+                if o.get("type") != "message":
+                    continue
+                m = o.get("message") or {}
+                if m.get("role") != "assistant":
+                    continue
+                u = m.get("usage")
+                if not u:
+                    continue
+                seen = True
+                if m.get("model"):
+                    model = m["model"]
+                tin += int(u.get("input") or 0)
+                tout += int(u.get("output") or 0)
+                tcr += int(u.get("cacheRead") or 0)
+                tcw += int(u.get("cacheWrite") or 0)
+                cost += float((u.get("cost") or {}).get("total") or 0.0)
+    except OSError:
+        return None
+    if not seen:
+        return None
+    return {
+        "model": model or "deepseek-v4-pro",
+        "input_tokens": tin,
+        "output_tokens": tout,
+        "cache_creation_tokens": tcw,
+        "cache_read_tokens": tcr,
+        "cost_reported": cost,
+        "duration_ms": None,
+    }
+
+
+def _session_cwd(path: str) -> Optional[str]:
+    """Read the header ``session`` line and return its ``cwd``, or None."""
+    try:
+        with open(path) as f:
+            for line in f:
+                line = line.strip()
+                if not line:
+                    continue
+                try:
+                    o = json.loads(line)
+                except json.JSONDecodeError:
+                    return None
+                if o.get("type") == "session":
+                    return o.get("cwd")
+                # session header is expected first; bail after first JSON line
+                return None
+    except OSError:
+        return None
+    return None
+
+
+def usage_from_session(
+    cwd: Optional[str] = None,
+    cycle_id: Optional[str] = None,
+    slug: Optional[str] = None,
+    base_dir: Optional[str] = None,
+) -> Optional[dict]:
+    """Recover a pi cycle's usage by reading its persisted session file(s).
+
+    Matching: scan ``<base>/*/*.jsonl`` and select files whose session
+    header ``cwd`` equals the target worktree path (authoritative).  When
+    ``cwd`` isn't given but ``cycle_id`` is, also accept files whose path
+    contains ``cycle-<cycle_id>`` (dir-name fallback).
+
+    Retries reuse the same worktree → multiple session files may match;
+    their usage is SUMMED (so token totals reflect wasted retry work too).
+
+    Returns the merged usage dict (tokens + model + ``cost_reported``), or
+    None when nothing matches / zero tokens (callers then skip writing,
+    preserving "n/a not fake zero").  The authoritative list cost is left to
+    the writer, which freezes it from the CNY price snapshot.
+    """
+    base = _sessions_base_dir(base_dir)
+    files = sorted(glob.glob(os.path.join(base, "*", "*.jsonl")))
+    if not files:
+        return None
+
+    matched = []
+    for path in files:
+        if cwd is not None and _session_cwd(path) == cwd:
+            matched.append(path)
+            continue
+        if cycle_id is not None and ("cycle-%s" % cycle_id) in path:
+            matched.append(path)
+
+    if not matched:
+        return None
+
+    agg = {
+        "model": None,
+        "input_tokens": 0,
+        "output_tokens": 0,
+        "cache_creation_tokens": 0,
+        "cache_read_tokens": 0,
+        "cost_reported": 0.0,
+        "duration_ms": None,
+    }
+    got = False
+    for path in matched:
+        s = _sum_session_file(path)
+        if s is None:
+            continue
+        got = True
+        agg["model"] = agg["model"] or s["model"]
+        agg["input_tokens"] += s["input_tokens"]
+        agg["output_tokens"] += s["output_tokens"]
+        agg["cache_creation_tokens"] += s["cache_creation_tokens"]
+        agg["cache_read_tokens"] += s["cache_read_tokens"]
+        agg["cost_reported"] += s["cost_reported"]
+
+    if not got:
+        return None
+    has_tokens = (
+        agg["input_tokens"] or agg["output_tokens"]
+        or agg["cache_creation_tokens"] or agg["cache_read_tokens"]
+    )
+    if not has_tokens:
+        return None
+    agg["model"] = agg["model"] or "deepseek-v4-pro"
+    return agg

package/lib/agent_usage/pi_emit.py

@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+"""
+pi_emit — write ONE authoritative usage event for a finished pi cycle.
+
+pi runs as ``pi -p`` (text mode): its stdout carries no token/cost summary,
+so loop-fmt's passthrough can only show progress, not capture usage. This
+thin CLI is invoked once by bin/roll after the agent phase (when ``$WT`` and
+``$CYCLE_ID`` are still in scope). It recovers the cycle's real usage from
+pi's persisted session files via ``pi.usage_from_session`` and appends a
+single ``stage=="usage"`` event to the loop events file.
+
+Exactly one event per cycle — the dashboard SUMS token fields across same-
+label usage events, so emitting once (here, post-cycle) instead of once per
+retry attempt (the old loop-fmt path) avoids ×N inflation.
+
+Cost is frozen at the active price snapshot in deepseek's native currency
+(CNY) via ``model_prices.compute_list_cost`` — the same convention claude
+uses (US-VIEW-014). pi's own ``cost.total`` (computed in USD) is kept as
+``cost_reported_usd`` for audit only. We never convert currencies; the
+dashboard already renders the right symbol from ``cost_currency``.
+
+When ``usage_from_session`` finds nothing (no session match, zero tokens),
+nothing is written — preserving "show n/a, not a fake zero".
+"""
+
+import argparse
+import importlib.util
+import json
+import os
+import sys
+from datetime import datetime, timezone
+
+_THIS_DIR = os.path.dirname(os.path.abspath(__file__))
+_LIB_DIR = os.path.dirname(_THIS_DIR)
+
+
+def _load_model_prices():
+    spec = importlib.util.spec_from_file_location(
+        "model_prices", os.path.join(_LIB_DIR, "model_prices.py")
+    )
+    mp = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(mp)
+    return mp
+
+
+def _load_pi():
+    spec = importlib.util.spec_from_file_location(
+        "agent_usage_pi", os.path.join(_THIS_DIR, "pi.py")
+    )
+    pi = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(pi)
+    return pi
+
+
+def build_event(cwd=None, cycle_id=None, slug=None, base_dir=None):
+    """Return the (line dict) usage event for a pi cycle, or None to skip.
+
+    None means no recoverable usage — caller writes nothing.
+    """
+    pi = _load_pi()
+    u = pi.usage_from_session(
+        cwd=cwd, cycle_id=cycle_id, slug=slug, base_dir=base_dir
+    )
+    if u is None:
+        return None
+
+    mp = _load_model_prices()
+    model = u.get("model") or "deepseek-v4-pro"
+    totals = {
+        "input_tokens": int(u.get("input_tokens") or 0),
+        "output_tokens": int(u.get("output_tokens") or 0),
+        "cache_creation_tokens": int(u.get("cache_creation_tokens") or 0),
+        "cache_read_tokens": int(u.get("cache_read_tokens") or 0),
+    }
+    cost_list = mp.compute_list_cost(model, **totals)
+    currency = mp.currency_for(model)
+
+    payload = {
+        "model": model,
+        "input_tokens": totals["input_tokens"],
+        "output_tokens": totals["output_tokens"],
+        "cache_creation_tokens": totals["cache_creation_tokens"],
+        "cache_read_tokens": totals["cache_read_tokens"],
+        # pi's own per-message cost.total summed, in USD — audit only.
+        "cost_reported_usd": u.get("cost_reported"),
+        "duration_ms": u.get("duration_ms"),
+        # Authoritative, frozen at snapshot in native currency (CNY).
+        "cost_list_usd": cost_list,
+        "cost_currency": currency,
+        "prices_version": getattr(mp, "VERSION", None),
+    }
+    return {
+        "ts": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
+        "stage": "usage",
+        "label": cycle_id,
+        "detail": payload,
+        "outcome": "ok",
+    }
+
+
+def _default_events_path(slug, shared):
+    base = shared or os.environ.get("LOOP_SHARED_ROOT") \
+        or os.path.expanduser("~/.shared/roll")
+    return os.path.join(base, "loop", "events-%s.ndjson" % slug)
+
+
+def main(argv=None):
+    ap = argparse.ArgumentParser(description="emit one pi usage event")
+    ap.add_argument("--cwd", help="cycle worktree path (authoritative match)")
+    ap.add_argument("--cycle", help="cycle id (label + dir-name fallback)")
+    ap.add_argument("--slug", help="project slug (events filename)")
+    ap.add_argument("--shared", help="shared root (for default events path)")
+    ap.add_argument("--events", help="explicit events file path (preferred)")
+    ap.add_argument("--base-dir", help="pi sessions root override (tests)")
+    args = ap.parse_args(argv)
+
+    event = build_event(
+        cwd=args.cwd, cycle_id=args.cycle, slug=args.slug, base_dir=args.base_dir
+    )
+    if event is None:
+        return 0  # nothing recoverable — write nothing (n/a, not fake zero)
+
+    evfile = args.events or _default_events_path(args.slug, args.shared)
+    try:
+        os.makedirs(os.path.dirname(evfile), exist_ok=True)
+        with open(evfile, "a") as f:
+            f.write(json.dumps(event) + "\n")
+    except OSError as e:
+        print("[pi_emit] failed to write %s: %s" % (evfile, e), file=sys.stderr)
+        return 1
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

package/lib/agent_usage/qwen.py

@@ -0,0 +1,128 @@
+"""
+qwen (Alibaba Qwen / dashscope CLI) agent usage extractor.
+
+Like openai, gemini and kimi (and unlike pi, which persists usage to session
+files), the Qwen / qwen-coder / dashscope CLI prints a token-usage summary to
+stdout at the end of a session.  So this plugin implements the standard
+``extract()`` registry contract: scrape the passthrough stdout lines for the
+usage / model lines.
+
+Recognised lines (case-insensitive, robust to thousands separators)::
+
+    Model: qwen-coder-plus
+    Tokens: input=15300 output=3120
+
+The dashscope "usage" / session-summary block is also accepted::
+
+    Input tokens:  15,300
+    Output tokens:  3,120
+    Total tokens:  18,420
+    model: qwen-max
+
+When an explicit USD cost line isn't present, cost is computed from
+``lib/model_prices.py`` (list price) so the dashboard never shows ``—``
+for a recognised qwen cycle.  Returns None if no usage line is found,
+so the caller falls back to the null payload (US-LOOP-010 compatible).
+"""
+
+import os
+import re
+import sys
+from typing import Optional
+
+# model_prices lives one level up (lib/), alongside this package.
+sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+try:
+    import model_prices
+except Exception:  # pragma: no cover - import guard
+    model_prices = None
+
+# Default model when the output omits an explicit model line.
+_DEFAULT_MODEL = "qwen-coder-plus"
+
+_MODEL_RE = re.compile(r"^\s*model\s*[:=]\s*([A-Za-z0-9][\w.\-]*)", re.IGNORECASE)
+_INPUT_RE = re.compile(r"input(?:\s+tokens)?\s*[:=]\s*([\d,]+)", re.IGNORECASE)
+_OUTPUT_RE = re.compile(r"output(?:\s+tokens)?\s*[:=]\s*([\d,]+)", re.IGNORECASE)
+_TOTAL_RE = re.compile(r"total(?:\s+tokens)?\s*[:=]\s*([\d,]+)", re.IGNORECASE)
+_COST_RE = re.compile(r"cost\s*[:=]?\s*\$?\s*([\d.]+)\s*(?:usd)?", re.IGNORECASE)
+
+
+def _to_int(s: str) -> int:
+    """Parse a token count string, tolerating thousands separators."""
+    return int(s.replace(",", ""))
+
+
+def extract(stdin_lines: list[str]) -> Optional[dict]:
+    """Parse Qwen CLI stdout and return a usage dict, or None.
+
+    Scans every line (the usage summary is at the tail but may be wrapped
+    by surrounding text) and accumulates the last seen model / token / cost
+    values.  Requires at least one of input/output/total tokens to be found;
+    otherwise returns None (caller falls back to null payload).
+    """
+    if not stdin_lines:
+        return None
+
+    model = None
+    tin = tout = ttotal = None
+    cost = None
+
+    for raw in stdin_lines:
+        line = raw.rstrip("\n")
+
+        m = _MODEL_RE.match(line)
+        if m:
+            model = m.group(1)
+
+        m = _INPUT_RE.search(line)
+        if m:
+            tin = _to_int(m.group(1))
+
+        m = _OUTPUT_RE.search(line)
+        if m:
+            tout = _to_int(m.group(1))
+
+        m = _TOTAL_RE.search(line)
+        if m:
+            ttotal = _to_int(m.group(1))
+
+        m = _COST_RE.search(line)
+        if m:
+            try:
+                cost = float(m.group(1))
+            except ValueError:
+                pass
+
+    # Require at least one token figure; otherwise this isn't a qwen cycle.
+    if tin is None and tout is None and ttotal is None:
+        return None
+    if tin is None and tout is None and ttotal is not None:
+        # No split available — attribute the whole total to input so the
+        # cycle is non-zero; output stays 0.
+        tin = ttotal
+        tout = 0
+    else:
+        tin = tin or 0
+        tout = tout or 0
+        if ttotal is not None and tin == 0 and tout == 0:
+            tin = ttotal
+
+    model = model or _DEFAULT_MODEL
+
+    if cost is None:
+        if model_prices is not None:
+            cost = model_prices.compute_list_cost(
+                model,
+                input_tokens=tin,
+                output_tokens=tout,
+            )
+        else:  # pragma: no cover - only when model_prices unimportable
+            cost = 0.0
+
+    return {
+        "model": model,
+        "input_tokens": tin,
+        "output_tokens": tout,
+        "cost_list_usd": cost,
+        "duration_ms": None,
+    }