DRU-1144 -- Count the tokens and cost of a Codex call's subagents

backend/druks/harnesses/codex.py

@@ -4,7 +4,7 @@
 import os
 import shlex
 from collections.abc import Collection, Iterator
-from dataclasses import dataclass
+from dataclasses import dataclass, fields
 from datetime import UTC, datetime
 from functools import lru_cache
 from pathlib import Path
@@ -127,6 +127,49 @@ def read_codex_cost_from_stdout(
     return _price(_normalize_usage(usage_raw), model=model)
 
 
+def read_codex_subagent_costs(
+    path: Path,
+    *,
+    stdout: bytes,
+    model: str | None,
+) -> Iterator[tuple[float, dict[str, Any]]]:
+    """Price each subagent that the call's main thread spawned, at the subagent's own model.
+
+    The wrapper writes up to three lines for each subagent session file: its
+    ``session_meta``, its last ``turn_context`` and its last ``token_count``.
+    The main thread's id comes from the ``thread.started`` event on stdout.
+    """
+    if not path.exists():
+        return
+    thread_id = None
+    for line in stdout.splitlines():
+        if b'"thread.started"' not in line:
+            continue
+        try:
+            event = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        if event["type"] == "thread.started":
+            thread_id = event["thread_id"]
+    is_child = False
+    subagent_model = model
+    for line in path.read_text(encoding="utf-8", errors="replace").splitlines():
+        try:
+            event = json.loads(line)
+        except json.JSONDecodeError:
+            continue
+        payload = event["payload"]
+        if event["type"] == "session_meta":
+            spawn = payload["source"]["subagent"]["thread_spawn"]
+            is_child = spawn["parent_thread_id"] == thread_id
+            subagent_model = model
+        elif event["type"] == "turn_context":
+            subagent_model = payload["model"]
+        elif is_child and payload["info"]:
+            usage = _normalize_usage(payload["info"]["total_token_usage"])
+            yield _price(usage, model=subagent_model)
+
+
 def _price(
     usage: _Usage,
     *,
@@ -427,6 +470,7 @@ def _build_codex_wrapper(
         in_vm_output = f"{in_vm_run_dir}/output.json"
         in_vm_marker = f"{in_vm_run_dir}/session.marker"
         in_vm_session = f"{in_vm_run_dir}/session.jsonl"
+        in_vm_subagents = f"{in_vm_run_dir}/subagents.jsonl"
         write_hook = ""
         if "after_agent" in hooks:
             command = f"sh {self.after_agent_hook_path(in_vm_run_dir)}"
@@ -468,6 +512,7 @@ def _build_codex_wrapper(
         codex_cmdline = " ".join(shlex.quote(a) for a in codex_argv)
         marker_q = shlex.quote(in_vm_marker)
         session_q = shlex.quote(in_vm_session)
+        subagents_q = shlex.quote(in_vm_subagents)
         # A subscription run writes its login before the command. Drukbox hands
         # the placeholder to the box only, so the box fills the file itself.
         login = ""
@@ -480,8 +525,12 @@ def _build_codex_wrapper(
             )
         # Codex runs against its real ``~/.codex`` because CODEX_HOME re-homes
         # all home-resolved state: auth, skills, and future features. A marker
-        # file plus ``-newer`` identifies the session JSONL. More than one match
-        # warns on stderr and copies nothing rather than mis-attributing cost.
+        # file plus ``-newer`` identifies the session JSONL. A subagent writes
+        # its own session file, so the wrapper sends only its first line, its
+        # last turn_context (the model) and its last token totals; parse keeps
+        # the ones this call spawned.
+        # More than one other match warns on stderr and copies nothing rather
+        # than mis-attributing cost.
         wrapper = (
             f"mkdir -p {shlex.quote(in_vm_run_dir)} && "
             f"printf %s {shlex.quote(schema_body)} > {shlex.quote(in_vm_schema)} && "
@@ -492,10 +541,16 @@ def _build_codex_wrapper(
             "trap : TERM && "
             f"{codex_cmdline}; "
             "ec=$?; "
-            f"sessions=$(find \"$HOME/.codex/sessions\" -type f -name '*.jsonl' "
-            f"-newer {marker_q} 2>/dev/null); "
-            "count=$(printf '%s\\n' \"$sessions\" | grep -c .); "
-            f'if [ "$count" -eq 1 ]; then cp "$sessions" {session_q}; '
+            "count=0; "
+            "for f in $(find \"$HOME/.codex/sessions\" -type f -name '*.jsonl' "
+            f"-newer {marker_q} 2>/dev/null); do "
+            'if head -n 1 "$f" | grep -q \'"thread_spawn":{\'; then '
+            'head -n 1 "$f"; '
+            'grep \'"type":"turn_context"\' "$f" | tail -n 1; '
+            'grep \'"type":"token_count"\' "$f" | tail -n 1; '
+            "else session=$f; count=$((count + 1)); fi; "
+            f"done > {subagents_q}; "
+            f'if [ "$count" -eq 1 ]; then cp "$session" {session_q}; '
             'elif [ "$count" -gt 1 ]; then '
             'echo "codex wrapper: $count session files newer than the marker;'
             ' refusing to guess which belongs to this run" >&2; fi; '
@@ -549,7 +604,7 @@ async def build_invocation(
             stdin=_with_final_message_note(prompt).encode("utf-8"),
             credentials=self._get_credentials(manifest),
             env=extra_env,
-            extra_artifact_filenames=("output.json", "session.jsonl"),
+            extra_artifact_filenames=("output.json", "session.jsonl", "subagents.jsonl"),
         )
 
     def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) -> Any:
@@ -560,6 +615,15 @@ def parse(self, result: HarnessRunResult, *, artifact_dir: Path, run_id: str) ->
                 call_dir / "session.jsonl",
                 model=self.model_id,
             )
+        if cost_metadata:
+            for subagent_cost, subagent_metadata in read_codex_subagent_costs(
+                call_dir / "subagents.jsonl",
+                stdout=result.stdout,
+                model=self.model_id,
+            ):
+                cost_usd = round(cost_usd + subagent_cost, 6)
+                for field in fields(_Usage):
+                    cost_metadata[field.name] += subagent_metadata[field.name]
         write_cost(call_dir, cost_usd=cost_usd, metadata=cost_metadata)
         self.check_returncode(result)
 

backend/tests/test_codex_cost.py

@@ -1,7 +1,10 @@
 import json
+import subprocess
 from pathlib import Path
 
+from druks.harnesses import codex
 from druks.harnesses.codex import CodexHarness, read_codex_cost_from_jsonl
+from druks.harnesses.datastructures import Manifest
 from druks.sandbox.datastructures import HarnessRunResult
 
 
@@ -139,3 +142,125 @@ def test_parse_falls_back_to_the_session_file_without_turn_completed(tmp_path: P
     cost = _parse_codex_call(tmp_path, stdout)
 
     assert cost["cost_usd"] == 11.25
+
+
+def _compact(event: dict) -> str:
+    return json.dumps(event, separators=(",", ":"))
+
+
+def _session_meta(thread_id: str, parent_thread_id: str) -> str:
+    spawn = {"parent_thread_id": parent_thread_id, "depth": 1}
+    payload = {"id": thread_id, "source": {"subagent": {"thread_spawn": spawn}}}
+    return _compact({"type": "session_meta", "payload": payload})
+
+
+def _turn_context(model: str) -> str:
+    return _compact({"type": "turn_context", "payload": {"model": model}})
+
+
+def _token_count(tokens: int) -> str:
+    usage = {"input_tokens": tokens, "cached_input_tokens": 0, "output_tokens": tokens}
+    payload = {"type": "token_count", "info": {"total_token_usage": usage}}
+    return _compact({"type": "event_msg", "payload": payload})
+
+
+_MAIN_TURN = (
+    b'{"type":"thread.started","thread_id":"main"}\n'
+    b'{"type":"turn.completed","usage":{"input_tokens":1000000,'
+    b'"cached_input_tokens":0,"output_tokens":1000000}}\n'
+)
+
+
+def test_parse_adds_a_subagent_priced_at_its_own_model(tmp_path: Path) -> None:
+    (tmp_path / "call").mkdir()
+    (tmp_path / "call" / "subagents.jsonl").write_text(
+        "\n".join(
+            [
+                _session_meta("child", parent_thread_id="main"),
+                _turn_context("gpt-5-mini"),
+                _token_count(1_000_000),
+            ]
+        )
+        + "\n"
+    )
+
+    cost = _parse_codex_call(tmp_path, _MAIN_TURN)
+
+    # $11.25 for the main thread at gpt-5.5, plus 1M input at $0.25 and
+    # 1M output at $2.00 for the subagent at gpt-5-mini.
+    assert cost["cost_usd"] == 13.5
+    assert cost["metadata"]["input_tokens"] == 2_000_000
+    assert cost["metadata"]["output_tokens"] == 2_000_000
+
+
+def test_parse_skips_a_subagent_of_another_call(tmp_path: Path) -> None:
+    (tmp_path / "call").mkdir()
+    (tmp_path / "call" / "subagents.jsonl").write_text(
+        "\n".join(
+            [
+                _session_meta("child", parent_thread_id="main"),
+                _turn_context("gpt-5.5"),
+                _token_count(1_000),
+                _session_meta("stranger", parent_thread_id="other"),
+                _turn_context("gpt-5.5"),
+                _token_count(1_000_000),
+            ]
+        )
+        + "\n"
+    )
+
+    cost = _parse_codex_call(tmp_path, _MAIN_TURN)
+
+    # $11.25 for the main thread, plus 1,000 input at $1.25 and 1,000 output
+    # at $10.00 per million for the subagent.
+    assert cost["cost_usd"] == 11.26125
+    assert cost["metadata"]["input_tokens"] == 1_001_000
+
+
+async def test_the_wrapper_sends_subagent_usage_beside_the_session(
+    tmp_path: Path, monkeypatch
+) -> None:
+    monkeypatch.setattr(codex, "get_runs_root", lambda _: str(tmp_path / "runs"))
+    sessions = tmp_path / "home" / ".codex" / "sessions" / "2026" / "11" / "10"
+    main = [_compact({"type": "session_meta", "payload": {"id": "main", "source": "exec"}})]
+    child = [
+        _session_meta("child", parent_thread_id="main"),
+        _turn_context("gpt-5.5"),
+        _token_count(1),
+        _compact({"type": "response_item", "payload": {"type": "message"}}),
+        _turn_context("gpt-5-mini"),
+        _token_count(2),
+    ]
+    fake_codex = tmp_path / "bin" / "codex"
+    fake_codex.parent.mkdir()
+    fake_codex.write_text(
+        "#!/bin/sh\n"
+        f"mkdir -p {sessions}\n"
+        f"cat > {sessions}/rollout-main.jsonl <<'EOF'\n" + "\n".join(main) + "\nEOF\n"
+        f"cat > {sessions}/rollout-child.jsonl <<'EOF'\n" + "\n".join(child) + "\nEOF\n"
+    )
+    fake_codex.chmod(0o755)
+    invocation = await CodexHarness(
+        model="openai/gpt-5.5", fast_mode=False, effort=None
+    ).build_invocation(
+        Manifest(model="", harness="", mcp_servers=(), skills=()),
+        prompt="hello",
+        schema={"type": "object"},
+        run_id="call",
+        ssh_username="exedev",
+    )
+
+    subprocess.run(
+        invocation.args,
+        input=b"",
+        env={"HOME": str(tmp_path / "home"), "PATH": f"{fake_codex.parent}:/usr/bin:/bin"},
+        check=True,
+    )
+
+    call_dir = tmp_path / "runs" / "call"
+    assert (call_dir / "session.jsonl").read_text().splitlines() == main
+    assert (call_dir / "subagents.jsonl").read_text().splitlines() == [
+        child[0],
+        child[4],
+        child[5],
+    ]

backend/tests/test_harness_reasoning_flags.py

@@ -170,7 +170,7 @@ async def test_codex_build_invocation_carries_every_flag():
     assert "CODEX_HOME" not in wrapper
     assert "druks-after-agent" not in wrapper
     assert inv.env == {"TOK": "secret"}
-    assert inv.extra_artifact_filenames == ("output.json", "session.jsonl")
+    assert inv.extra_artifact_filenames == ("output.json", "session.jsonl", "subagents.jsonl")
 
 
 async def test_codex_blocks_the_first_stop_of_a_hooked_call():