import json
import sys
import unittest
from pathlib import Path
HERE = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(HERE))
sys.path.insert(0, str(HERE / "tests"))
import test_cli # noqa: E402
from test_usage import entry # noqa: E402
from wflib import usage # noqa: E402
SID = "99999999-2222-3333-4444-555555555555"
USER = json.dumps({"type": "user", "timestamp": "2026-10-05T09:00:00Z", "message": {"role": "user", "content": "hi"}})
def transcript(*sizes):
"""One request per size (inp 10, cw 1000, rest cache read), each streamed as 2 entries."""
lines = [USER]
for n, size in enumerate(sizes):
e = entry(f"m{n}", f"r{n}", "claude-opus-5-5", f"2026-10-05T09:0{n}:01Z", inp=10, cw5=1000, cr=size - 1010)
lines += [e, e, USER]
return "\n".join(lines) + "\n"
class ContextTokens(unittest.TestCase):
def test_last_request_prompt_size(self):
self.assertEqual(usage.context_tokens(transcript(50_000, 152_000)), 152_000)
def test_skips_synthetic_and_sidechain(self):
text = transcript(40_000)
synth = json.loads(entry("s", "rs", "<synthetic>", "2026-10-05T10:00:00Z", inp=1))
side = json.loads(entry("x", "rx", "claude-opus-5-5", "2026-10-05T10:00:01Z", cr=999_000))
side["isSidechain"] = True
text += json.dumps(synth) + "\n" + json.dumps(side) + "\n"
self.assertEqual(usage.context_tokens(text), 40_000)
def test_partial_first_line_and_empty(self):
self.assertEqual(usage.context_tokens('ens": 5}}}\n' + transcript(30_000)), 30_000)
self.assertIsNone(usage.context_tokens(USER + "\n"))
def test_hint_line(self):
self.assertEqual(usage.ctx_hint(152_000, 100_000),
"context ~152k tokens (> 100k): ask the owner to /clear, then continue "
"(subagent: ignore, this is the main session)")
self.assertIsNone(usage.ctx_hint(99_999, 100_000))
self.assertIsNone(usage.ctx_hint(None, 100_000))
self.assertIsNone(usage.ctx_hint(500_000, 0))
class HintCli(test_cli.Cli):
def setUp(self):
super().setUp()
self.home = self.root.parent / "claude-home"
(self.home / "projects" / "-work-demo").mkdir(parents=True)
self.jsonl = self.home / "projects" / "-work-demo" / f"{SID}.jsonl"
def env(self, sid=SID):
return {"CLAUDE_CONFIG_DIR": str(self.home), "CLAUDE_CODE_SESSION_ID": sid}
HINT = "context ~152k tokens (> 100k)"
def test_done_prints_hint_when_big(self):
self.jsonl.write_text(transcript(152_000))
out = self.ok("done", "t-one", "-m", "ok", env=self.env())
self.assertIn(self.HINT, out.splitlines()[-1])
def test_next_prints_hint_when_big(self):
self.jsonl.write_text(transcript(152_000))
out = self.ok("next", "--as", "opus", env=self.env())
self.assertIn(self.HINT, out.splitlines()[-1])
def test_silent_when_small_or_missing(self):
self.jsonl.write_text(transcript(60_000))
self.assertNotIn("context ~", self.ok("next", "--as", "opus", env=self.env()))
self.assertNotIn("context ~", self.ok("next", "--as", "opus", env=self.env(sid="")))
self.assertNotIn("context ~", self.ok("next", "--as", "opus", env=self.env(sid="nope")))
def test_threshold_from_toml(self):
(self.root / "workflow.toml").write_text(self.toml + "ctx_hint = 200000\n")
self.jsonl.write_text(transcript(152_000))
self.assertNotIn("context ~", self.ok("next", "--as", "opus", env=self.env()))
(self.root / "workflow.toml").write_text(self.toml + "ctx_hint = 0\n")
self.jsonl.write_text(transcript(900_000))
self.assertNotIn("context ~", self.ok("next", "--as", "opus", env=self.env()))
def test_bad_threshold(self):
(self.root / "workflow.toml").write_text(self.toml + 'ctx_hint = "big"\n')
self.assertIn("'ctx_hint' must be a number", self.fails("list"))
def test_reads_only_the_tail(self):
self.jsonl.write_text(transcript(152_000) + (USER + "\n") * 40_000 + transcript(160_000))
self.assertIn("context ~160k tokens", self.ok("next", "--as", "opus", env=self.env()))
if __name__ == "__main__":
unittest.main()