diff --git a/backend/druks/harnesses/codex.py b/backend/druks/harnesses/codex.py index d8b2c2eb..88e0f298 100644 --- a/backend/druks/harnesses/codex.py +++ b/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) diff --git a/backend/tests/test_codex_cost.py b/backend/tests/test_codex_cost.py index 75c2bb87..ac574a86 100644 --- a/backend/tests/test_codex_cost.py +++ b/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], + ] diff --git a/backend/tests/test_harness_reasoning_flags.py b/backend/tests/test_harness_reasoning_flags.py index 012354a7..8f6051aa 100644 --- a/backend/tests/test_harness_reasoning_flags.py +++ b/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():