path: root/tests/test_usage.py

import json
import os
import subprocess
import sys
import tempfile
import unittest
from pathlib import Path

HERE = Path(__file__).resolve().parent.parent
WF = HERE / "wf.py"
sys.path.insert(0, str(HERE))

from wflib import usage  # noqa: E402


def entry(mid, req, model, ts, inp=0, cw5=0, cw1h=0, cr=0, out=0, type="assistant", stop=None, content=None,
          uuid=None):
    u = {"input_tokens": inp, "cache_creation_input_tokens": cw5 + cw1h, "cache_read_input_tokens": cr,
         "output_tokens": out}
    if cw1h:
        u["cache_creation"] = {"ephemeral_5m_input_tokens": cw5, "ephemeral_1h_input_tokens": cw1h}
    d = {"type": type, "requestId": req, "timestamp": ts,
         "message": {"id": mid, "model": model, "role": "assistant", "usage": u}}
    if stop:
        d["message"]["stop_reason"] = stop
    if content is not None:
        d["message"]["content"] = content
    if uuid:
        d["uuid"] = uuid
    return json.dumps(d)


# Subagent transcripts: stop_reason always null, output_tokens = message_start placeholder.
# e1: thinking (signature 1800 chars) + text 1000 chars + tool_use (input json 100 chars), placeholder 8.
#   by hand: 0.3*(1800-800) + 0.3*1000 + (30 + 0.44*100) = 300 + 300 + 74 = 674
# e2: same tool_use again but stop_reason set (final): reported 50 is trusted.
# e3: no content, placeholder 9 kept (nothing to estimate).
TOOL = {"type": "tool_use", "id": "t1", "name": "Bash", "input": {"a": "y" * 91}}  # json.dumps → 100 chars
SUB = "\n".join([
    entry("e1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:00Z", out=8, uuid="u1",
          content=[{"type": "thinking", "thinking": "", "signature": "z" * 1800}]),
    entry("e1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:01Z", out=8, uuid="u2",
          content=[{"type": "text", "text": "x" * 1000}]),
    entry("e1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:01Z", out=8, uuid="u2",
          content=[{"type": "text", "text": "x" * 1000}]),  # duplicated line: counted once
    entry("e1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:02Z", out=8, uuid="u3", content=[TOOL]),
    entry("e2", "q2", "claude-sonnet-5-5", "2026-10-04T10:01:00Z", out=50, stop="tool_use", content=[TOOL]),
    entry("e3", "q3", "claude-sonnet-5-5", "2026-10-04T10:02:00Z", out=9),
]) + "\n"


# main session, opus 5.5: m1 streamed as 2 entries (same usage), m2 as 2 entries (out 4000 then 10000).
# by hand: m1 = 1000*4 + 100000*5 + 2000*20 = 544000 µ$; m2 = 50000*8 + 1e6*0.2 + 10000*20 = 800000 µ$
MAIN = "\n".join([
    json.dumps({"type": "user", "timestamp": "2026-10-04T09:00:00Z", "message": {"role": "user", "content": "hi"}}),
    entry("m1", "r1", "claude-opus-5-5", "2026-10-04T09:00:01Z", inp=1000, cw5=100000, out=2000),
    entry("m1", "r1", "claude-opus-5-5", "2026-10-04T09:00:02Z", inp=1000, cw5=100000, out=2000),
    "not json",
    json.dumps({"type": "attachment", "timestamp": "2026-10-04T09:00:03Z"}),
    entry("m2", "r2", "claude-opus-5-5", "2026-10-04T09:01:00Z", cw1h=50000, cr=1000000, out=4000),
    entry("m2", "r2", "claude-opus-5-5", "2026-10-04T09:01:05Z", cw1h=50000, cr=1000000, out=10000),
    entry("m9", "r9", "<synthetic>", "2026-10-04T09:02:00Z", out=5),
]) + "\n"

# subagent, sonnet 5.5: s1 = 10*2 + 20000*2.5 + 30000*0.2 + 500*10 = 61020 µ$; s2 = 50000*0.2 + 1500*10 = 25000 µ$
AGENT = "\n".join([
    entry("s1", "q1", "claude-sonnet-5-5", "2026-10-04T10:00:00Z", inp=10, cw5=20000, cr=30000, out=500),
    entry("s2", "q2", "claude-sonnet-5-5", "2026-10-04T12:00:00Z", cr=50000, out=1500),
    entry("s3", "q3", "claude-mystery-1", "2026-10-04T12:30:00Z", inp=7, out=3),
]) + "\n"

SID = "11111111-2222-3333-4444-555555555555"


class Parse(unittest.TestCase):
    def test_dedupes_streamed_entries_and_takes_last_output(self):
        got = usage.parse(MAIN)
        self.assertEqual(list(got), ["claude-opus-5-5"])
        u = got["claude-opus-5-5"]
        self.assertEqual((u.turns, u.inp, u.cw5, u.cw1h, u.cr, u.out), (2, 1000, 100000, 50000, 1000000, 12000))
        self.assertEqual(u.cw, 150000)
        self.assertAlmostEqual(usage.cost("claude-opus-5-5", u), 1.344)

    def test_models_kept_apart_unknown_price_none(self):
        got = usage.parse(AGENT)
        s = got["claude-sonnet-5-5"]
        self.assertEqual((s.turns, s.cr, s.out), (2, 80000, 2000))
        self.assertAlmostEqual(usage.cost("claude-sonnet-5-5", s), 0.08602)
        self.assertIsNone(usage.cost("claude-mystery-1", got["claude-mystery-1"]))

    def test_since_until_filter_on_utc_timestamp(self):
        s = usage.parse(AGENT, since="2026-10-04T11:00", until="2026-10-04T12:10")["claude-sonnet-5-5"]
        self.assertEqual((s.turns, s.cr), (1, 50000))
        self.assertAlmostEqual(usage.cost("claude-sonnet-5-5", s), 0.025)

    def test_dated_model_id_priced_by_prefix(self):
        u = usage.Usage(turns=1, inp=1000000, out=1000000)
        self.assertAlmostEqual(usage.cost("claude-haiku-4-5-20251001", u), 6.0)
        self.assertAlmostEqual(usage.cost("claude-opus-5", u), 30.0)

    def test_subagent_output_estimated_from_content(self):
        u = usage.parse(SUB)["claude-sonnet-5-5"]
        self.assertEqual((u.turns, u.out, u.est), (3, 674 + 50 + 9, 1))

    def test_log_line_marks_estimate(self):
        line = usage.log_line("T", "p", "t-x", "1h", "done", "a1", usage.parse(SUB))
        self.assertIn(" out=733 est=1 usd=", line)
        self.assertNotIn(" est=", usage.log_line("T", "p", "t-x", "1h", "done", "a1", usage.parse(AGENT)))

    def test_log_line_lane_and_model(self):
        by = usage.parse(SUB)
        self.assertIn(" lane=fast model=sonnet effort=", usage.log_line("T", "p", "t-x", "1h", "done", "a1", by, lane="fast"))
        self.assertIn(" lane=sonnet model=sonnet effort=", usage.log_line("T", "p", "t-x", "1h", "done", "a1", by))

    def test_report_lane_model_key(self):
        es = [{"lane": "fast", "model": "sonnet", "outcome": "done", "usd": 1.0, "turns": 10, "effort": "1h"},
              {"lane": "fast", "model": "opus", "outcome": "done", "usd": 2.0, "turns": 20, "effort": "1h"},
              {"lane": "opus", "outcome": "done", "usd": 3.0, "turns": 30, "effort": "1h"}]
        self.assertEqual(usage.report(es, ("1h",)), [
            ("fast/opus", "all", 1, 1, 0, 2.0, 2.0, 2.0, 20, None), ("fast/opus", "1h", 1, 1, 0, 2.0, 2.0, 2.0, 20, None),
            ("fast/sonnet", "all", 1, 1, 0, 1.0, 1.0, 1.0, 10, None), ("fast/sonnet", "1h", 1, 1, 0, 1.0, 1.0, 1.0, 10, None),
            ("opus", "all", 1, 1, 0, 3.0, 3.0, 3.0, 30, None), ("opus", "1h", 1, 1, 0, 3.0, 3.0, 3.0, 30, None)])

    def test_fmt(self):
        self.assertEqual([usage.fmt(n) for n in (999, 1500, 8948413)], ["999", "1.5k", "8.95M"])


def tcall(mid, ctx, name, inp, tid):
    return json.dumps({"type": "assistant", "timestamp": "2026-10-04T10:00:00Z", "message": {
        "id": mid, "model": "claude-sonnet-5-5", "usage": {"input_tokens": 10, "cache_read_input_tokens": ctx - 10},
        "content": [{"type": "tool_use", "id": tid, "name": name, "input": inp}]}})


def tres(tid, text):
    return json.dumps({"type": "user", "message": {"role": "user", "content": [
        {"type": "tool_result", "tool_use_id": tid, "content": text}]}})


# agent A: ctx 1000 Read(400 chars=100 tok), 2000 grep(800=200), 3000 Edit, 4000 git commit.
#   exp 3000 / mut 3000 / book 4000 of 10000; first edit after 2 calls, pre 3000 = 30%, context +2000
EXA = "\n".join([
    tcall("a1", 1000, "Read", {"file_path": "/w/proj/.worktrees/x/src/a.py"}, "t1"), tres("t1", "x" * 400),
    tcall("a2", 2000, "Bash", {"command": "grep -n foo src/a.py 2>/dev/null"}, "t2"), tres("t2", "y" * 800),
    tcall("a3", 3000, "Edit", {"file_path": "/w/proj/src/a.py"}, "t3"), tres("t3", "ok"),
    tcall("a4", 4000, "Bash", {"command": "git commit -m x"}, "t4"), tres("t4", "done"),
]) + "\n"
# agent B: 4 exploring calls of 1000 each, reads src/a.py (100 tok) + sed -n of b.py (400 chars)
EXB = "\n".join([
    tcall("b1", 1000, "Read", {"file_path": "/w/proj/src/a.py"}, "u1"), tres("u1", "x" * 400),
    tcall("b2", 1000, "Bash", {"command": "sed -n 1,9p src/b.py"}, "u2"), tres("u2", "z" * 400),
    tcall("b3", 1000, "Bash", {"command": "ls"}, "u3"), tres("u3", "q"),
    tcall("b4", 1000, "Glob", {}, "u4"), tres("u4", "q"),
]) + "\n"


class Explore(unittest.TestCase):
    def test_one_agent(self):
        e = usage.explore(EXA, root="/w/proj")
        self.assertEqual((e.calls, e.first_edit, e.pre, e.ctx_growth), (4, 2, 3000, 2000))
        self.assertEqual(e.cost, {"exp": 3000, "mut": 3000, "book": 4000})
        self.assertEqual(e.results, {"Read": 100, "grep": 200, "Edit": 0, "other": 1})
        self.assertEqual(e.files, {"src/a.py": 300})

    def test_call_kind_book_is_writes_only(self):
        kind = lambda c: usage.call_kind({"name": "Bash", "input": {"command": c}})
        for c in ("git worktree list", "git branch --list 'fast/*'", "git branch --show-current", "git status --short"):
            self.assertEqual(kind(c), "exp", c)
        for c in ("git worktree add .w/x -b x master", "git worktree remove .w/x", "git branch -d x",
                  "python3 /projects/public/workflow/wf.py finish t-x -m ok", "git switch -c x"):
            self.assertEqual(kind(c), "book", c)

    def test_since_drops_early_calls(self):
        e = usage.explore(EXA.replace("10:00:00Z", "09:00:00Z", 2), since="2026-10-04T10")
        self.assertEqual(e.calls, 2)

    def test_report(self):
        a, b = usage.explore(EXA, root="/w/proj"), usage.explore(EXB, root="/w/proj")
        r = usage.explore_report([("A", a), ("B", b)], min_pre=4)
        self.assertEqual(r["rows"], [("A", 4, 2, 0.3, 0.3), ("B", 4, 4, 1.0, 0.0)])
        self.assertEqual((r["total"], r["cost"]["exp"]), (14000, 7000))
        self.assertEqual((r["pre_n"], r["pre_share"], r["pre_calls"], r["pre_ctx"]), (1, 0.3, 2, 2000))
        self.assertEqual(r["results"]["sed-cat"], 100)
        self.assertEqual(r["files"], [("src/a.py", 2, 400)])
        self.assertEqual(usage.explore_report([("A", a)], min_calls=5)["rows"], [])


class Cli(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.TemporaryDirectory()
        self.home = Path(self.tmp.name)
        proj = self.home / "projects" / "-work-demo"
        sub = proj / SID / "subagents"
        sub.mkdir(parents=True)
        (proj / f"{SID}.jsonl").write_text(MAIN)
        (sub / "agent-abc123.jsonl").write_text(AGENT)
        (sub / "agent-abc123.meta.json").write_text(json.dumps({"description": "wf-worker sonnet t-x"}))

    def tearDown(self):
        self.tmp.cleanup()

    def wf(self, *args, sid="", cwd=None):
        env = {**os.environ, "CLAUDE_CONFIG_DIR": str(self.home), "CLAUDE_CODE_SESSION_ID": sid}
        r = subprocess.run([sys.executable, str(WF), "usage", *args], capture_output=True, text=True,
                           cwd=cwd or self.tmp.name, env=env, timeout=30)
        return r.returncode, r.stdout, r.stderr

    def test_explore_skips_main_and_short_agents(self):
        code, out, err = self.wf("--session", SID[:8], "--explore")
        self.assertEqual((code, err), (0, ""))
        self.assertEqual(out.strip(), "no subagent with >= 4 calls")

    def test_session_rows_and_total(self):
        code, out, err = self.wf("--session", SID[:8])
        self.assertEqual((code, err), (0, ""))
        rows = [l.split() for l in out.strip().split("\n")[1:]]
        self.assertEqual(rows[0], ["main", "opus-5-5", "2", "1.0k", "150.0k", "1.00M", "12.0k", "1.34"])
        self.assertEqual(rows[1], ["abc123", "wf-worker", "sonnet", "t-x",
                                   "sonnet-5-5", "2", "10", "20.0k", "80.0k", "2.0k", "0.09"])
        self.assertEqual(rows[2][-7:], ["mystery-1", "1", "7", "0", "0", "3", "?"])
        self.assertEqual(rows[3], ["total", "5", "1.0k", "170.0k", "1.08M", "14.0k", "1.43+"])

    def test_session_from_env(self):
        code, out, _ = self.wf(sid=SID)
        self.assertEqual(code, 0)
        self.assertIn("abc123", out)

    def test_agent_only(self):
        code, out, err = self.wf("--agent", "agent-abc123", "--since", "2026-10-04T11:00")
        self.assertEqual((code, err), (0, ""))
        lines = out.strip().split("\n")
        self.assertNotIn("main", out)
        self.assertEqual(lines[1].split()[-7:], ["sonnet-5-5", "1", "0", "0", "50.0k", "1.5k", "0.03"])

    def test_log_appends_one_key_value_line(self):
        root = self.home / "demo"
        root.mkdir()
        (root / "workflow.toml").write_text("format = 1\n")
        code, out, err = self.wf("--agent", "abc123", "--log", "t-x", "done", "--effort", "1h", cwd=root)
        self.assertEqual((code, err), (0, ""))
        code, out, err = self.wf("--agent", "abc123", "--log", "t-x", "handback", cwd=root)
        log = (root / "out" / "wf-cost.log").read_text().split("\n")
        self.assertEqual(len(log), 3)  # two lines + trailing newline
        # lane = model with most turns (sonnet 2 vs mystery 1); usd = known models only
        self.assertRegex(log[0], r"^\d{4}-\d\d-\d\dT\d\d:\d\d:\d\dZ project=demo task=t-x lane=sonnet model=sonnet effort=1h "
                                 r"outcome=done turns=3 in=17 cw=20000 cr=80000 out=2003 usd=0\.0860 agent=abc123$")
        self.assertIn(" effort=- outcome=handback ", log[1])
        self.assertEqual(out, log[1] + "\n")

    def test_estimated_out_marked_with_tilde(self):
        other = self.home / "projects" / "-work-demo" / "other-session" / "subagents"
        other.mkdir(parents=True)
        (other / "agent-est1.jsonl").write_text(SUB)
        code, out, err = self.wf("--agent", "est1")
        self.assertEqual((code, err), (0, ""))
        self.assertEqual(out.strip().split("\n")[1].split()[-3:], ["0", "~733", "0.01"])

    def test_log_needs_agent(self):
        code, out, err = self.wf("--log", "t-x", "done")
        self.assertEqual((code, err), (2, "wf: --log needs --agent\n"))

    def test_log_effort_must_be_an_estimate(self):
        code, out, err = self.wf("--agent", "abc123", "--log", "t-x", "done", "--effort", "medium")
        self.assertEqual(code, 2)
        self.assertIn("invalid choice: 'medium'", err)

    def test_missing_agent_one_line_error(self):
        code, out, err = self.wf("--agent", "nope")
        self.assertEqual((code, out), (1, ""))
        self.assertEqual(err, "wf: no transcript for agent nope\n")


LOG_A = """\
2026-10-01T10:00:00Z project=a task=t-1 lane=sonnet effort=1h outcome=done turns=10 usd=0.50
2026-10-01T11:00:00Z project=a task=t-2 lane=sonnet effort=1h outcome=handback turns=30 usd=1.50
garbage line
"""
LOG_B = """\
2026-10-02T10:00:00Z project=b task=t-3 lane=sonnet effort=<1h outcome=done turns=20 usd=0.40
2026-10-02T11:00:00Z project=b task=t-4 lane=opus effort=5h outcome=done turns=80 usd=3.00
2026-10-03T11:00:00Z project=b task=t-5 lane=opus effort=1h outcome=done turns=40 usd=2.00
"""


class Report(unittest.TestCase):
    def test_rows_by_lane_then_effort(self):
        entries = usage.parse_log(LOG_A + LOG_B)
        self.assertEqual(len(entries), 5)
        rows = usage.report(entries, ("<1h", "1h", "5h"))
        # lane, effort, n, done, other, median $, total $, $/done, median turns — all by hand
        self.assertEqual(rows, [
            ("opus", "all", 2, 2, 0, 2.5, 5.0, 2.5, 60, None),
            ("opus", "1h", 1, 1, 0, 2.0, 2.0, 2.0, 40, None),
            ("opus", "5h", 1, 1, 0, 3.0, 3.0, 3.0, 80, None),
            ("sonnet", "all", 3, 2, 1, 0.5, 2.4, 1.2, 20, None),
            ("sonnet", "<1h", 1, 1, 0, 0.4, 0.4, 0.4, 20, None),
            ("sonnet", "1h", 2, 1, 1, 1.0, 2.0, 2.0, 20, None),
        ])

    def test_duration(self):
        line = usage.log_line("T", "p", "t-x", "1h", "done", "a1", usage.parse(SUB), dur=125)
        self.assertIn(" dur=125 agent=a1", line)
        self.assertNotIn("dur=", usage.log_line("T", "p", "t-x", "1h", "done", "a1", usage.parse(SUB)))
        log = ("2026-10-04T10:00:00Z project=a task=t-6 lane=fast effort=<1h outcome=done turns=1 usd=1 dur=100\n"
               "2026-10-04T10:00:00Z project=a task=t-7 lane=fast effort=<1h outcome=done turns=1 usd=1 dur=300\n"
               "2026-10-04T10:00:00Z project=a task=t-8 lane=fast effort=<1h outcome=done turns=1 usd=1\n")
        self.assertEqual(usage.report(usage.parse_log(log), ("<1h",))[0][9], 200)

    def test_done_gate_red_counts_as_done(self):
        log = ("2026-10-04T10:00:00Z project=a task=t-6 lane=fast effort=<1h outcome=done+gate-red turns=12 usd=1.00\n"
               "2026-10-04T11:00:00Z project=a task=t-7 lane=fast effort=<1h outcome=donex turns=8 usd=0.50\n")
        rows = usage.report(usage.parse_log(log), ("<1h",))
        self.assertEqual(rows[0], ("fast", "all", 2, 1, 1, 0.75, 1.5, 1.5, 10, None))

    def test_since_and_no_done(self):
        entries = usage.parse_log(LOG_A + LOG_B, since="2026-10-01T10:30")
        rows = usage.report(entries, ("1h",))
        self.assertIn(("sonnet", "all", 2, 1, 1, 0.95, 1.9, 1.9, 25, None), rows)
        rows = usage.report(usage.parse_log(LOG_A, since="2026-10-01T10:30"), ("1h",))
        self.assertEqual(rows[0], ("sonnet", "all", 1, 0, 1, 1.5, 1.5, None, 30, None))

    def test_cli_reads_every_project(self):
        with tempfile.TemporaryDirectory() as tmp:
            for name, text in (("a", LOG_A), ("b", LOG_B)):
                (Path(tmp) / name / "out").mkdir(parents=True)
                (Path(tmp) / name / "workflow.toml").write_text("format = 1\n")
                (Path(tmp) / name / "out" / "wf-cost.log").write_text(text)
            r = subprocess.run([sys.executable, str(WF), "usage", "--report"], capture_output=True, text=True,
                               cwd=tmp, env={**os.environ, "WF_ROOT": tmp}, timeout=30)
        self.assertEqual((r.returncode, r.stderr), (0, ""))
        lines = [l.split() for l in r.stdout.strip().split("\n")]
        self.assertEqual(lines[0], ["lane", "effort", "n", "done", "other", "med$", "total$", "$/done", "med_turns", "med_dur"])
        self.assertEqual(lines[1], ["opus", "all", "2", "2", "0", "2.50", "5.00", "2.50", "60", "-"])
        self.assertEqual(lines[6], ["sonnet", "1h", "2", "1", "1", "1.00", "2.00", "2.00", "20", "-"])


if __name__ == "__main__":
    unittest.main()