diff --git a/skills/review-switch-codex/scripts/tui_review_bridge.py b/skills/review-switch-codex/scripts/tui_review_bridge.py index ccb22da..493e1e9 100644 --- a/skills/review-switch-codex/scripts/tui_review_bridge.py +++ b/skills/review-switch-codex/scripts/tui_review_bridge.py @@ -75,6 +75,25 @@ # The vendor half of the lane this bridge reviews in. The model half is whatever # the lineage was pinned to, so the lane needs no argument of its own. REVIEW_VENDOR = "codex" +# The four disjoint counters a `session-cost` line carries, in the machine log's +# own spelling. +COUNTERS = ("input", "output", "cache_read", "cache_creation") +# Where Codex keeps the rollout of every thread it has run, one JSONL file per +# thread under a dated tree, each filename ending in the thread's own id. +CODEX_HOME_ENV_VAR = "CODEX_HOME" +CODEX_HOME_DEFAULT = ".codex" +ROLLOUT_DIRECTORY = "sessions" +ROLLOUT_GLOB = "**/*.jsonl" +# The counters a rollout's cumulative `total_token_usage` must carry for it to be +# read at all. `reasoning_output_tokens` is not among them: it counts a subset of +# the output tokens rather than a fifth pot, and nothing here needs it. +# A `token_count` event that carried no usage object: held apart from "no count at all", because +# what it means is that this thread's last word on what it spent is unreadable. +MALFORMED_COUNT = object() +ROLLOUT_USAGE_FIELDS = ( + "input_tokens", "cached_input_tokens", "cache_write_input_tokens", + "output_tokens", "total_tokens", +) class AppServerError(RuntimeError): @@ -104,6 +123,9 @@ def __init__(self, log, ticket, model, vendor=REVIEW_VENDOR): # field verbatim after collapsing whitespace, so this is the spelling the # dashboard shows. self.lane = " ".join(part for part in (vendor, model) if part) + # A review leaves exactly one cost line per call, whether it came back + # with figures or with the reason there are none. + self.costed = False def write(self, state): """Append one end of this review; returns nothing, and fails at nothing. @@ -127,6 +149,180 @@ def write(self, state): except OSError: pass + def cost(self, session, model, counters, detail): + """Append this review's one `session-cost` line; returns nothing, and fails at nothing. + + Lane-tagged, so a review's spend is told apart from the implementing + child's, and written through the log's own writer, which holds the + figures-or-diagnosis rule: `counters` are the four disjoint totals, or + `None` with `detail` saying why nobody could tell. + + Accounting must never fail a review, so every way this can go wrong is + swallowed here, exactly as the `review` pair's write is. + """ + if not self.log or not self.ticket: + return + self.costed = True + command = [ + sys.executable, str(MACHINE_LOG), "--log", str(self.log), + "session-cost", "--ticket", str(self.ticket), + "--executor", REVIEW_VENDOR, "--model", model or "", "--lane", self.lane, + ] + if session: + command += ["--session", str(session)] + if counters is None: + command += ["--detail", detail] + else: + for name in COUNTERS: + command += [f"--{name.replace('_', '-')}-tokens", str(counters[name])] + command += ["--total-tokens", str(sum(counters.values()))] + try: + subprocess.run(command, capture_output=True, text=True, check=False) + except OSError: + pass + + +def codex_sessions_root(environment=None): + """The directory Codex writes its rollouts under, wherever this machine keeps it.""" + environment = os.environ if environment is None else environment + home = environment.get(CODEX_HOME_ENV_VAR) + root = pathlib.Path(home) if home else pathlib.Path.home() / CODEX_HOME_DEFAULT + return root / ROLLOUT_DIRECTORY + + +def find_rollout(root, thread_id): + """The rollout file this thread wrote, or `None` where the tree holds none. + + The id glob is the lookup: Codex names every rollout for the thread it + belongs to, so the filename ending in the id the bridge already holds is the + file, without reading a line of any other. + """ + try: + found = sorted( + path for path in root.glob(ROLLOUT_GLOB) if path.stem.endswith(thread_id) + ) + except OSError: + return None + return found[-1] if found else None + + +def rollout_counters(usage): + """Codex's cumulative counters as the machine log's four disjoint ones, or `None`. + + Codex counts its cached tokens inside `input_tokens`, so they come back out + here, and its reasoning tokens inside `output_tokens`, where they stay: the + four that come out are disjoint, and their sum is the total Codex itself + reported. + + `None` is every way the source does not hold together — a counter missing, + not a count, or negative; more cached tokens than input tokens; a mapping + that does not add up to the reported total. A figure invented out of a + contradiction is worse than the diagnosis that says the rollout could not be + read, because nothing downstream can tell it apart from a measurement. + """ + counts = {} + for name in ROLLOUT_USAGE_FIELDS: + value = usage.get(name) + if not isinstance(value, int) or isinstance(value, bool) or value < 0: + return None + counts[name] = value + cached = counts["cached_input_tokens"] + if cached > counts["input_tokens"]: + return None + counters = { + "input": counts["input_tokens"] - cached, + "output": counts["output_tokens"], + "cache_read": cached, + "cache_creation": counts["cache_write_input_tokens"], + } + if sum(counters.values()) != counts["total_tokens"]: + return None + return counters + + +def harvest_rollout(root, thread_id): + """What one review thread spent: `(counters, model, detail)`, read off its rollout. + + A read of a file the lane already wrote, so it costs no model token and drives + nothing. `counters` is `None` when the figures could not be read, and `detail` + then says why, which is the diagnosis the `session-cost` line carries instead. + + The **last** cumulative count is the answer: every turn appends one, and a + resumed round two appends to the same file, so one read covers a whole review + however many rounds it had. + + Only the file's last line is forgiven for not parsing: the pane this reads + behind is still attached, so that one may be half written and is unbilled + either way. A line that does not parse with more of the rollout written after + it is a hole in the history, and is diagnosed rather than quietly skipped. + """ + path = find_rollout(root, thread_id) + if path is None: + return None, None, f"no Codex rollout under {root} ends in {thread_id}" + reported = None + model = None + unparsed_last = False + try: + with path.open(encoding="utf-8") as stream: + for line in stream: + if unparsed_last: + return None, model, f"{path} carries a line that does not parse" + try: + record = json.loads(line) + except json.JSONDecodeError: + unparsed_last = True + continue + if not isinstance(record, dict): + unparsed_last = True + continue + payload = record.get("payload") + if not isinstance(payload, dict): + continue + if record.get("type") == "turn_context" and payload.get("model"): + model = payload["model"] + if payload.get("type") == "token_count": + info = payload.get("info") + usage = info.get("total_token_usage") if isinstance(info, dict) else None + # The count this event carries, whatever shape it is in: a `token_count` + # whose usage is missing or is not an object is not an event to skip back + # past, because the count before it is no longer this thread's last word. + reported = usage if isinstance(usage, dict) else MALFORMED_COUNT + except OSError as error: + return None, model, f"{path} could not be read: {error.strerror}" + if reported is None: + return None, model, f"{path} reports no token count" + if reported is MALFORMED_COUNT: + return None, model, f"{path}'s last token count reports no usage" + counters = rollout_counters(reported) + if counters is None: + return None, model, f"{path} reports a token count that does not hold together" + return counters, model, None + + +def resumed_thread_id(args): + """The thread a resumed call was about to review, or `None` when it names none. + + A failed round still spent tokens in that thread, and the id is what keeps its rollout out of + a same-vendor child's figures, so the diagnosis row names it wherever the record survives. + """ + if not getattr(args, "resume_session", None): + return None + try: + return SessionStore().read(args.resume_session).get("threadId") + except (OSError, RuntimeError, ValueError): + return None + + +def record_review_cost(review, thread_id): + """Leave this review's one `session-cost` line, harvested from its own rollout. + + The model recorded is the one the thread resolved to and nothing else: the + alias the caller asked for is already on the row, in the lane, and writing it + into the model field would dress a guess up as a measurement. + """ + counters, model, detail = harvest_rollout(codex_sessions_root(), thread_id) + review.cost(thread_id, model, counters, detail) + @dataclasses.dataclass(frozen=True) class InvocationOwner: @@ -1240,7 +1436,7 @@ async def run_recovered_review(args, owner, store): ) -async def run_bridge(args): +async def run_bridge(args, review=None): if not pathlib.Path(args.cwd).is_dir(): raise RuntimeError(f"Working directory does not exist: {args.cwd}") owner = resolve_owner(args) @@ -1269,6 +1465,10 @@ async def run_bridge(args): cleanup_required = should_cleanup_session(turn, final_message) if cleanup_required: cleanup_session(state, store) + if review is not None: + # Once the turn has settled, so the rollout's last cumulative count is + # the whole of what this review spent. + record_review_cost(review, state["threadId"]) output = { "status": turn.get("status"), "reviewSessionId": state["reviewSessionId"], @@ -1373,7 +1573,7 @@ def main(): review = ReviewEvent(args.machine_log, args.ticket, args.model) review.write("running") try: - return asyncio.run(run_bridge(args)) + return asyncio.run(run_bridge(args, review)) except NoLiveSessionError as error: print(str(error), file=sys.stderr) return NO_LIVE_SESSION_EXIT @@ -1381,6 +1581,14 @@ def main(): print(str(error), file=sys.stderr) return 1 finally: + # A review that never came back read no counters, and one that is left + # out of the log entirely is the blind spot this event exists to close: + # the diagnosis goes in where the figures cannot. + if not review.costed: + review.cost( + resumed_thread_id(args), None, None, + "this review returned no result to read a cost from", + ) review.write("returned") diff --git a/skills/review-switch-codex/tests/test_tui_review_bridge.py b/skills/review-switch-codex/tests/test_tui_review_bridge.py index 0cf72c6..63452bc 100644 --- a/skills/review-switch-codex/tests/test_tui_review_bridge.py +++ b/skills/review-switch-codex/tests/test_tui_review_bridge.py @@ -1005,6 +1005,405 @@ def test_round_two_writes_its_own_pair_for_the_same_ticket(self): self.assertEqual({record["ticket"] for record in records}, {self.TICKET}) + +def token_count(usage): + """One `token_count` event, in the shape a Codex rollout writes it.""" + return { + "timestamp": "2026-08-18T04:00:00.000Z", + "type": "event_msg", + "payload": {"type": "token_count", "info": {"total_token_usage": usage}}, + } + + +def turn_context(model): + """One `turn_context` line, which is where a rollout says which model the turn resolved to.""" + return { + "timestamp": "2026-08-18T04:00:00.000Z", + "type": "turn_context", + "payload": {"model": model, "cwd": "/workspace"}, + } + + +def write_rollout(root, thread_id, records, day="2026/08/18"): + """One rollout file, under the dated tree and named the way Codex names it.""" + directory = root / day + directory.mkdir(parents=True, exist_ok=True) + path = directory / f"rollout-2026-08-18T04-00-00-{thread_id}.jsonl" + path.write_text( + "".join(json.dumps(record) + "\n" for record in records), encoding="utf-8" + ) + return path + + +# One round's cumulative counters, as Codex reports them: `input_tokens` counts the cached ones +# inside itself, and `reasoning_output_tokens` counts a subset of `output_tokens`. +ROUND_ONE_USAGE = { + "input_tokens": 200000, + "cached_input_tokens": 180000, + "cache_write_input_tokens": 0, + "output_tokens": 9000, + "reasoning_output_tokens": 6000, + "total_tokens": 209000, +} +ROUND_TWO_USAGE = { + "input_tokens": 350000, + "cached_input_tokens": 310000, + "cache_write_input_tokens": 0, + "output_tokens": 15000, + "reasoning_output_tokens": 9000, + "total_tokens": 365000, +} +# The same figures after the mapping the spec pins: cache-read is the cached count, input is what +# is left of the input count once the cached ones are out of it, cache-creation is the write +# count, and output keeps its reasoning tokens inside it rather than beside them. +ROUND_ONE_COUNTERS = { + "input": 20000, "output": 9000, "cache_read": 180000, "cache_creation": 0, +} +ROUND_TWO_COUNTERS = { + "input": 40000, "output": 15000, "cache_read": 310000, "cache_creation": 0, +} +RESOLVED_MODEL = "gpt-5.6-sol" + + +class RolloutHarvestTests(unittest.TestCase): + """Reading what a Codex review spent out of the rollout its thread id names. + + A pure read of a file the lane already wrote: no model token is spent to obtain it, and every + way it can fail answers with the diagnosis that goes in the `session-cost` line instead. + """ + + def setUp(self): + self.bridge = load_bridge() + self.work = tempfile.TemporaryDirectory() + self.addCleanup(self.work.cleanup) + self.sessions = pathlib.Path(self.work.name) / "sessions" + + def harvest(self, thread_id="thread-ticket-13"): + return self.bridge.harvest_rollout(self.sessions, thread_id) + + def test_a_rollout_is_found_by_the_thread_id_its_filename_ends_in(self): + write_rollout( + self.sessions, "thread-ticket-13", + [turn_context(RESOLVED_MODEL), token_count(ROUND_ONE_USAGE)], + ) + write_rollout( + self.sessions, "another-thread", + [turn_context("gpt-5.6-luna"), token_count(ROUND_TWO_USAGE)], + ) + + counters, model, detail = self.harvest() + + self.assertIsNone(detail) + self.assertEqual(counters, ROUND_ONE_COUNTERS) + self.assertEqual(model, RESOLVED_MODEL) + + def test_the_counters_are_de_overlapped_so_they_sum_to_the_total_codex_reported(self): + write_rollout( + self.sessions, "thread-ticket-13", + [turn_context(RESOLVED_MODEL), token_count(ROUND_ONE_USAGE)], + ) + + counters, _model, _detail = self.harvest() + + self.assertEqual(sum(counters.values()), ROUND_ONE_USAGE["total_tokens"]) + + def test_a_resumed_round_is_read_as_one_review_from_the_last_cumulative_count(self): + """Round two appends to the same rollout, and its counters already cover round one.""" + write_rollout( + self.sessions, "thread-ticket-13", + [ + turn_context(RESOLVED_MODEL), + token_count(ROUND_ONE_USAGE), + turn_context(RESOLVED_MODEL), + token_count(ROUND_TWO_USAGE), + ], + ) + + counters, _model, detail = self.harvest() + + self.assertIsNone(detail) + self.assertEqual(counters, ROUND_TWO_COUNTERS) + + def test_a_rollout_that_counted_nothing_is_diagnosed_rather_than_billed_at_zero(self): + path = write_rollout( + self.sessions, "thread-ticket-13", [turn_context(RESOLVED_MODEL)] + ) + + counters, _model, detail = self.harvest() + + self.assertIsNone(counters) + self.assertIn(str(path), detail) + + def test_a_thread_with_no_rollout_on_disk_is_diagnosed(self): + write_rollout(self.sessions, "another-thread", [token_count(ROUND_ONE_USAGE)]) + + counters, model, detail = self.harvest() + + self.assertIsNone(counters) + self.assertIsNone(model) + self.assertIn("thread-ticket-13", detail) + self.assertIn(str(self.sessions), detail) + + def test_a_count_that_contradicts_itself_is_diagnosed_rather_than_invented(self): + """More cached tokens than input tokens is not a figure to clamp into shape.""" + path = write_rollout( + self.sessions, "thread-ticket-13", + [ + turn_context(RESOLVED_MODEL), + token_count({ + "input_tokens": 100, + "cached_input_tokens": 400, + "cache_write_input_tokens": 0, + "output_tokens": 50, + "reasoning_output_tokens": 10, + "total_tokens": 150, + }), + ], + ) + + counters, _model, detail = self.harvest() + + self.assertIsNone(counters) + self.assertIn(str(path), detail) + + def test_a_count_that_does_not_add_up_to_its_own_total_is_diagnosed(self): + """The four mapped counters are the source's own total, or they are not its counters.""" + write_rollout( + self.sessions, "thread-ticket-13", + [ + turn_context(RESOLVED_MODEL), + token_count(dict(ROUND_ONE_USAGE, total_tokens=209001)), + ], + ) + + counters, _model, detail = self.harvest() + + self.assertIsNone(counters) + self.assertIsNotNone(detail) + + def test_a_count_missing_one_of_its_figures_is_diagnosed_rather_than_read_as_zero(self): + missing = dict(ROUND_ONE_USAGE) + missing.pop("cache_write_input_tokens") + write_rollout( + self.sessions, "thread-ticket-13", + [turn_context(RESOLVED_MODEL), token_count(missing)], + ) + + counters, _model, detail = self.harvest() + + self.assertIsNone(counters) + self.assertIsNotNone(detail) + + def test_a_line_that_does_not_parse_mid_rollout_is_diagnosed(self): + """A hole in the history is not a line to skip: what came after it is unaccounted for.""" + path = write_rollout( + self.sessions, "thread-ticket-13", + [turn_context(RESOLVED_MODEL), token_count(ROUND_ONE_USAGE)], + ) + path.write_text( + path.read_text(encoding="utf-8") + "{not json at all\n" + + json.dumps(token_count(ROUND_TWO_USAGE)) + "\n", + encoding="utf-8", + ) + + counters, _model, detail = self.harvest() + + self.assertIsNone(counters) + self.assertIn(str(path), detail) + + def test_an_unreadable_last_count_is_diagnosed_rather_than_answered_by_an_older_one(self): + """The last count is the thread's word on what it spent; an older one is not a stand-in.""" + path = write_rollout( + self.sessions, "thread-ticket-13", + [ + turn_context(RESOLVED_MODEL), + token_count(ROUND_ONE_USAGE), + { + "timestamp": "2026-08-18T04:00:00.000Z", + "type": "event_msg", + "payload": {"type": "token_count", "info": {}}, + }, + ], + ) + + counters, _model, detail = self.harvest() + + self.assertIsNone(counters) + self.assertIn(str(path), detail) + + def test_a_half_written_last_line_does_not_lose_the_count_before_it(self): + """The pane may still be writing, so a truncated tail is not a lost review.""" + path = write_rollout( + self.sessions, "thread-ticket-13", + [turn_context(RESOLVED_MODEL), token_count(ROUND_ONE_USAGE)], + ) + with path.open("a", encoding="utf-8") as stream: + stream.write('{"timestamp": "2026-08-18T04:0') + + counters, _model, detail = self.harvest() + + self.assertIsNone(detail) + self.assertEqual(counters, ROUND_ONE_COUNTERS) + + + +class SessionCostTests(FakePaneTestCase): + """The one `session-cost` call a Codex review makes beside its `review` pair. + + Same seam as the `review` pair, tested the same way: the writer belongs to + the consumer that configured `--machine-log`, so what this lane is + responsible for is the argv it hands over, and that argv is what these + tests pin. + """ + + TICKET = "26" + MODEL_ALIAS = "sol" + + def setUp(self): + super().setUp() + self.machine_log = self.root / "run" / "log.jsonl" + writer = self.root / "machine_log.py" + writer.write_text(STUB_WRITER, encoding="utf-8") + self.calls_file = writer.with_name("calls.jsonl") + self.enter(mock.patch.object(self.bridge, "MACHINE_LOG", writer)) + self.codex_home = self.root / "codex-home" + os.environ["CODEX_HOME"] = str(self.codex_home) + self.addCleanup(os.environ.pop, "CODEX_HOME", None) + + def write_rollout(self, records, thread_id=None): + return write_rollout( + self.codex_home / "sessions", + thread_id or self.codex.thread_id, + records, + ) + + def main(self, *arguments): + argv = ["tui_review_bridge.py", "--cwd", str(self.worktree), *arguments] + stdout = io.StringIO() + with mock.patch.object(sys, "argv", argv): + with contextlib.redirect_stdout(stdout): + return self.bridge.main() + + def review_argv(self, *arguments): + return ( + "HEAD", + "--model", self.MODEL_ALIAS, + "--timeout", "5", + "--startup-timeout", "5", + "--machine-log", str(self.machine_log), + "--ticket", self.TICKET, + *arguments, + ) + + def costs(self): + """Every `session-cost` call the bridge made, read back as the flags it passed.""" + if not self.calls_file.exists(): + return [] + records = [] + for line in self.calls_file.read_text(encoding="utf-8").splitlines(): + if not line.strip(): + continue + arguments = json.loads(line) + if "session-cost" not in arguments: + continue + records.append({ + flag.lstrip("-"): value + for flag, value in zip(arguments, arguments[1:]) + if flag.startswith("--") + }) + return records + + def test_a_review_records_what_its_thread_spent(self): + self.write_rollout([turn_context(RESOLVED_MODEL), token_count(ROUND_ONE_USAGE)]) + self.codex.finish("one spec finding") + + code = self.main(*self.review_argv()) + + self.assertEqual(code, 0) + costs = self.costs() + self.assertEqual(len(costs), 1, costs) + entry = costs[0] + self.assertEqual(entry["ticket"], self.TICKET) + self.assertEqual(entry["executor"], "codex") + self.assertEqual(entry["lane"], f"codex {self.MODEL_ALIAS}") + self.assertEqual(entry["session"], self.codex.thread_id) + # The model the thread resolved to, not the alias the caller asked for. + self.assertEqual(entry["model"], RESOLVED_MODEL) + # Every counter reaches the writer as the string a command line carries. + self.assertEqual(entry["input-tokens"], str(ROUND_ONE_COUNTERS["input"])) + self.assertEqual(entry["output-tokens"], str(ROUND_ONE_COUNTERS["output"])) + self.assertEqual(entry["cache-read-tokens"], str(ROUND_ONE_COUNTERS["cache_read"])) + self.assertEqual( + entry["cache-creation-tokens"], str(ROUND_ONE_COUNTERS["cache_creation"]) + ) + self.assertEqual(entry["total-tokens"], str(ROUND_ONE_USAGE["total_tokens"])) + self.assertNotIn("detail", entry) + + def test_a_review_whose_rollout_is_missing_records_the_diagnosis(self): + self.codex.finish("no findings") + + code = self.main(*self.review_argv()) + + self.assertEqual(code, 0) + entry = self.costs()[0] + self.assertIn(self.codex.thread_id, entry["detail"]) + self.assertNotIn("total-tokens", entry) + + def test_a_review_that_never_came_back_records_the_diagnosis(self): + """The turn never finishes, so there is nothing to read — and the call still goes out.""" + self.write_rollout([turn_context(RESOLVED_MODEL), token_count(ROUND_ONE_USAGE)]) + + code = self.main( + "HEAD", "--model", self.MODEL_ALIAS, "--timeout", "0.2", + "--startup-timeout", "5", "--machine-log", str(self.machine_log), + "--ticket", self.TICKET, + ) + + self.assertEqual(code, 1) + entry = self.costs()[0] + self.assertEqual(entry["lane"], f"codex {self.MODEL_ALIAS}") + self.assertTrue(entry["detail"]) + self.assertNotIn("total-tokens", entry) + + def test_a_failed_second_round_still_names_the_thread_it_spent_in(self): + """Naming it is what keeps a same-vendor review out of the reviewed child's figures.""" + self.write_rollout([turn_context(RESOLVED_MODEL), token_count(ROUND_ONE_USAGE)]) + self.codex.finish("round one findings") + self.main(*self.review_argv()) + session = self.stored_session()["reviewSessionId"] + self.codex.status = "in_progress" + + code = self.main( + "HEAD", "--model", self.MODEL_ALIAS, "--timeout", "0.2", + "--startup-timeout", "5", "--machine-log", str(self.machine_log), + "--ticket", self.TICKET, "--resume-session", session, + ) + + self.assertEqual(code, 1) + self.assertEqual(self.costs()[-1]["session"], self.codex.thread_id) + + def test_a_rollout_that_names_no_model_records_no_model_rather_than_the_alias(self): + """The alias is already on the row, in the lane; the model field is a measurement.""" + self.write_rollout([token_count(ROUND_ONE_USAGE)]) + self.codex.finish("no findings") + + code = self.main(*self.review_argv()) + + self.assertEqual(code, 0) + self.assertEqual(self.costs()[0]["model"], "") + + def test_a_review_with_no_log_configured_records_no_cost(self): + self.write_rollout([turn_context(RESOLVED_MODEL), token_count(ROUND_ONE_USAGE)]) + self.codex.finish("no findings") + + code = self.main( + "HEAD", "--model", self.MODEL_ALIAS, "--timeout", "5", "--startup-timeout", "5" + ) + + self.assertEqual(code, 0) + self.assertEqual(self.costs(), []) + + class RecoveryParserTests(unittest.TestCase): def setUp(self): self.bridge = load_bridge()