From 0f2417437440ae5d9d0be6a95d185f67ab642a9e Mon Sep 17 00:00:00 2001 From: Aaron Breckenridge Date: Thu, 10 Sep 2026 08:55:07 -0500 Subject: [PATCH] Track LLM cost in Datadog LLM Obs for the BiggiePockets review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Each LLM span in the review trace now carries what the pass actually used and, where it is known, what it actually cost, so spend is attributed per model on the same trace as the review-quality metrics. No rate table lives in this repository. A list price committed to a file goes stale silently and would be reported with the same confidence as a real one. Instead each span carries whichever of the two things Datadog needs to price it: token counts for a model in its catalog, and a `total_cost` metric for one it does not carry. That cost is the amount OpenRouter reports charging, from the `cost` field it returns on every response. Datadog's catalog keys on the bare model and the provider that originated it, so the OpenRouter slug is split ("openai/gpt-5.6-sol" -> gpt-5.6-sol + openai) and the routing is recorded as a gateway tag. Without that split nothing would match and no model would be priced at all. The two passes record usage differently. Claude Code writes a `usage` object to its execution output. Codex writes running token counters to a session rollout on its own runner and its action exposes no usage output, so the Codex job reads that rollout and hands the totals to the reporting job through the existing handoff. Claude Code's own total_cost_usd is ignored: it is computed against Anthropic's list prices while these passes are billed by OpenRouter for a non-Anthropic model. The Datadog site and ml_app move to the DD_SITE and DD_LLMOBS_ML_APP org-level variables, so no Datadog endpoint or org identifier is committed here. Both fall back to public defaults. Only usage objects are read — never message content, transcripts, prompts, or diffs. Reporting stays best-effort and cannot fail a review. Co-Authored-By: Claude Opus 5 (1M context) --- .github/workflows/biggiepockets-review.yml | 59 +++++- .gitignore | 2 + README.md | 44 ++++ scripts/llm-usage.py | 226 +++++++++++++++++++++ tests/test_llm_usage.py | 180 ++++++++++++++++ 5 files changed, 504 insertions(+), 7 deletions(-) create mode 100644 .gitignore create mode 100755 scripts/llm-usage.py create mode 100644 tests/test_llm_usage.py diff --git a/.github/workflows/biggiepockets-review.yml b/.github/workflows/biggiepockets-review.yml index 02adde7..f6fc533 100644 --- a/.github/workflows/biggiepockets-review.yml +++ b/.github/workflows/biggiepockets-review.yml @@ -9,7 +9,13 @@ name: BiggiePockets Code Review (reusable) # Secrets are supplied by the caller via `secrets: inherit`: # OPENROUTER_API_KEY, JIRA_EMAIL, JIRA_API_TOKEN, # BIGGIEPOCKETS_PAT, DATADOG_API_KEY (optional; enables review-quality -# metrics in Datadog LLM Obs). +# metrics and cost tracking in Datadog LLM Obs). +# +# Cost tracking: each LLM span in the review trace carries the token counts and, +# where OpenRouter reported one, the cost the pass was charged, so Datadog attributes +# spend per model on the same trace as the quality metrics. Two optional org-level +# variables say where that lands: DD_SITE and DD_LLMOBS_ML_APP. No Datadog endpoint +# or org identifier is committed here. See README.md. # # Prompt versioning: every review-stage prompt lives in this repo's prompts/ # registry (BiggerPockets/.github) as a versioned template, plus one shared rule @@ -246,12 +252,17 @@ jobs: run: | codex_end_ns=$(date +%s%N) echo "CODEX_END_NS=$codex_end_ns" >> "$GITHUB_ENV" + # Codex writes its token counters to a session rollout on this runner, and + # the action exposes no usage output, so read them here and hand them to the + # job that reports the trace. One line, because the handoff is read as env. + codex_usage=$(python3 registry/scripts/llm-usage.py codex "${CODEX_MODEL:-unspecified}" | tr -d '\n') || codex_usage="" { echo "BASE_REF=$BASE_REF" echo "REVIEW_START_NS=$REVIEW_START_NS" echo "CODEX_START_NS=$CODEX_START_NS" echo "CODEX_END_NS=$codex_end_ns" echo "CODEX_STATUS=${{ steps.codex.outcome }}" + echo "CODEX_USAGE_JSON=$codex_usage" } > review-context.env touch codex-findings.md @@ -418,6 +429,11 @@ jobs: continue-on-error: true env: DD_API_KEY: ${{ secrets.DATADOG_API_KEY }} + # Site and ml_app come from org-level variables so no Datadog endpoint or + # org identifier is committed here; both fall back to public defaults. + DD_SITE: ${{ vars.DD_SITE }} + DD_LLMOBS_ML_APP: ${{ vars.DD_LLMOBS_ML_APP }} + EXECUTION_FILE: ${{ steps.claude.outputs.execution_file }} CODEX_PROMPT_TEMPLATE: ${{ steps.resolve.outputs.codex_prompt_template }} ARM_PROMPT_TEMPLATE: ${{ steps.resolve.outputs.arm_prompt_template }} run: | @@ -508,6 +524,30 @@ jobs: codex_model_tag="${CODEX_MODEL:-unspecified}" claude_model_tag="${CLAUDE_MODEL:-unspecified}" + dd_site="${DD_SITE:-datadoghq.com}" + ml_app="${DD_LLMOBS_ML_APP:-biggiepockets-review}" + + # Cost tracking. Each pass carries its own usage as span metrics plus a + # model identity split into the bare name and originating provider that + # Datadog's pricing catalog keys on ("openai/gpt-5.6-sol" -> gpt-5.6-sol + + # openai); the OpenRouter gateway is kept as a tag instead. Datadog then + # prices a catalogued model from its token counts, and takes the reported + # `total_cost` at face value for one it doesn't carry. Both numbers come + # from what the tools recorded — OpenRouter reports what it charged, and + # llm-usage.py holds no rate table of its own. + # + # Only usage objects are read, never message content. The Codex pass's + # counters travel from its own job in CODEX_USAGE_JSON, since the action + # exposes no usage output and its rollout stays on that runner. + # Fall back to an empty object if the helper can't run at all, so a cost + # problem can't cost us the quality metrics in the same payload. + no_fields='{"model_name":"unspecified","model_provider":"unspecified","metrics":{}}' + codex_fields="${CODEX_USAGE_JSON:-}" + claude_fields=$(python3 registry/scripts/llm-usage.py claude "$claude_model_tag" "${EXECUTION_FILE:-}") || claude_fields="$no_fields" + usable='has("model_name") and has("model_provider") and has("metrics")' + echo "$codex_fields" | jq -e "$usable" >/dev/null 2>&1 || codex_fields="$no_fields" + echo "$claude_fields" | jq -e "$usable" >/dev/null 2>&1 || claude_fields="$no_fields" + run_id="$GITHUB_REPOSITORY-pr$PR-${{ github.run_id }}" trace_id=$(openssl rand -hex 16) @@ -534,6 +574,8 @@ jobs: --arg codex_status "$codex_status" \ --arg claude_status "$claude_status" \ --arg root_status "$root_status" \ + --argjson codex_fields "$codex_fields" \ + --argjson claude_fields "$claude_fields" \ --arg codex_findings_excerpt "$codex_findings_excerpt" \ --arg summary_excerpt "$summary_excerpt" \ --argjson codex_prompt "$codex_prompt_obj" \ @@ -561,10 +603,11 @@ jobs: start_ns: $codex_start_ns, duration: $codex_duration, status: $codex_status, + metrics: $codex_fields.metrics, meta: { kind: "llm", - model_name: $codex_model, - model_provider: "openrouter", + model_name: $codex_fields.model_name, + model_provider: $codex_fields.model_provider, input: { messages: [{role: "user", content: "Review pr:\($pr) against ticket.json/pr.diff"}], prompt: $codex_prompt @@ -580,10 +623,11 @@ jobs: start_ns: $claude_start_ns, duration: $claude_duration, status: $claude_status, + metrics: $claude_fields.metrics, meta: { kind: "llm", - model_name: $claude_model, - model_provider: "openrouter", + model_name: $claude_fields.model_name, + model_provider: $claude_fields.model_provider, input: { messages: [{role: "user", content: $codex_findings_excerpt}], prompt: $arm_prompt @@ -594,7 +638,7 @@ jobs: ]') payload=$(jq -n \ - --arg ml_app "biggiepockets-review" \ + --arg ml_app "$ml_app" \ --arg repo "$GITHUB_REPOSITORY" \ --arg pr "$PR" \ --arg run_id "$run_id" \ @@ -629,6 +673,7 @@ jobs: "control_arm:\($control_arm)", "assignment_bucket:\($assignment_bucket)", "experiment_split_percent:\($split_percent)", + "gateway:openrouter", "codex_model:\($codex_model)", "claude_model:\($claude_model)", "codex_findings_lines:\($codex_findings_lines)", @@ -641,7 +686,7 @@ jobs: }') http_code=$(curl -sS -o /tmp/dd_response.json -w '%{http_code}' \ - -X POST "https://api.datadoghq.com/api/intake/llm-obs/v1/trace/spans" \ + -X POST "https://api.${dd_site}/api/intake/llm-obs/v1/trace/spans" \ -H "Content-Type: application/json" \ -H "DD-API-KEY: $DD_API_KEY" \ -d "$payload") || http_code="000" diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..7a60b85 --- /dev/null +++ b/.gitignore @@ -0,0 +1,2 @@ +__pycache__/ +*.pyc diff --git a/README.md b/README.md index 07c51b3..f0db58f 100644 --- a/README.md +++ b/README.md @@ -100,6 +100,50 @@ run and post normally without it, but no metrics are reported. The exact secret names each step expects are visible in the `env:` and `with:` blocks of [`.github/workflows/biggiepockets-review.yml`](.github/workflows/biggiepockets-review.yml). +#### Cost tracking (Datadog LLM Observability) + +`secrets.DATADOG_API_KEY` also turns on cost tracking. Every LLM span in the review +trace carries its token usage — input, output, total, and cached-read counts — under +`metrics`, and Datadog prices the span from its own model pricing catalog. Cost is +therefore attributed per pass and per model on the same trace as the quality metrics, +and shows up in LLM Observability's spend views without a separate report. + +For the catalog to recognise a model, a span has to name it the way Datadog does: the +bare model and the provider that originated it. The workflow runs everything through +OpenRouter, whose slugs look like `openai/gpt-5.6-sol`, so `scripts/llm-usage.py` +splits the slug into `model_name: gpt-5.6-sol` / `model_provider: openai` and records +the routing as a `gateway:openrouter` tag. + +Cost itself is never computed here, and there is no rate table in the repository: a +list price committed to a file goes stale silently and would be reported with the same +confidence as a real one. Instead each span carries whichever of the two things +Datadog needs. For a model in the catalog, token counts are enough. For one it does +not carry, the span reports a `total_cost` metric — the amount OpenRouter says it +charged, taken from the `cost` field it returns on every response, where that figure +survives into what the tool wrote to disk. + +The two passes record their usage differently. Claude Code writes a `usage` object to +its execution output. Codex writes running token counters to a session rollout on its +own runner, and since its action exposes no usage output, the workflow reads that +rollout in the Codex job and hands the totals to the reporting job. In both cases only +usage objects are read — never message content, transcripts, prompts, or diffs. Claude +Code's own `total_cost_usd` is ignored: it is computed against Anthropic's list prices, +while these passes are billed by OpenRouter for a non-Anthropic model. + +Two **organization-level variables** (`vars`, not secrets — **Settings → Secrets and +variables → Actions → Variables** at the org level) configure where the trace lands. +Both are optional, and nothing here is committed to the repository: + +- `DD_SITE` — the Datadog site to report to (e.g. `datadoghq.eu`, `us5.datadoghq.com`). + Defaults to the public `datadoghq.com`. Set this to the organization's actual site; + a private or internal Datadog hostname belongs in this variable and nowhere else. +- `DD_LLMOBS_ML_APP` — the LLM Obs `ml_app` the review trace is grouped under. + Defaults to `biggiepockets-review`. + +Cost tracking is best-effort and never fails a review. With no `DATADOG_API_KEY` the +whole reporting step is skipped, and a missing execution file, an absent rollout, or +malformed usage data degrades to fewer metrics on the span. + #### 4. Set workflow permissions The reusable workflow's jobs need `pull-requests: read` and `id-token: write` (OIDC diff --git a/scripts/llm-usage.py b/scripts/llm-usage.py new file mode 100755 index 0000000..f4ccce8 --- /dev/null +++ b/scripts/llm-usage.py @@ -0,0 +1,226 @@ +#!/usr/bin/env python3 +"""Emit the usage and cost fields for one review pass's Datadog LLM Obs span. + +Datadog prices an LLM span two ways. When `model_name`/`model_provider` name a model +in its pricing catalog, token counts on the span are enough — it costs the span +itself. When they don't, it needs the money: a `total_cost` metric it takes at face +value. This script produces whichever of the two the pass can support, from numbers +the tools themselves recorded. It holds no rate table. A list price copied into this +file would go stale silently and be reported with the same confidence as a real one. + +Where the numbers come from: + +- Both passes reach models through OpenRouter, which reports what it charged in the + `cost` field of every response's usage object. Where that figure survives into what + the tool wrote to disk, it is the authoritative cost for the pass and is emitted as + `total_cost` — a real charge, not an estimate. +- The Codex pass writes token counters to its session rollout. Its model is in the + catalog, so those counts are enough for Datadog to price it. +- Claude Code's own `total_cost_usd` is deliberately ignored. It is computed against + Anthropic's list prices, and these passes are billed by OpenRouter for a non- + Anthropic model, so it describes a bill nobody was sent. + +The workflow names models the way OpenRouter routes them ("openai/gpt-5.6-sol"). +Datadog's catalog keys on the bare model and its originating provider, so +`split_model` splits that into ("gpt-5.6-sol", "openai"); the fact that the call was +routed through OpenRouter is preserved separately as a `gateway` tag. + +Only counts, costs, and model identifiers pass through here. Message content, +transcripts, prompts, diffs, and responses are never read or emitted. + +Usage: llm-usage.py claude [execution-file] + llm-usage.py codex [rollout-dir] +Prints a JSON object on stdout for the workflow's jq to splice into a span: + {"model_name": ..., "model_provider": ..., "gateway": ..., "metrics": {...}} +`metrics` carries only what is actually known; it is `{}` when nothing is. +Never raises and always exits 0 — cost telemetry must not fail a code review. +""" +import json +import os +import sys + +# Datadog's catalog keys on the originating provider, not the gateway a call was +# routed through. Slugs the workflow uses map to that provider by their first segment. +GATEWAY = "openrouter" + +DEFAULT_ROLLOUT_DIR = os.path.expanduser("~/.codex/sessions") + +# Claude Code's usage object, in Anthropic's naming, mapped to Datadog's metric names. +CLAUDE_USAGE_FIELDS = { + "input_tokens": "input_tokens", + "output_tokens": "output_tokens", + "cache_read_input_tokens": "cache_read_input_tokens", + "cache_creation_input_tokens": "cache_write_input_tokens", +} + +# Codex's rollout counters, in its own naming, mapped to Datadog's metric names. +# `cached_input_tokens` is a subset of `input_tokens`, matching Datadog's split of +# input into non-cached and cache-read. +CODEX_USAGE_FIELDS = { + "input_tokens": "input_tokens", + "output_tokens": "output_tokens", + "cached_input_tokens": "cache_read_input_tokens", +} + + +def split_model(slug): + """"openai/gpt-5.6-sol" -> ("gpt-5.6-sol", "openai"). A slug with no provider + prefix keeps its name and reports an unspecified provider.""" + slug = (slug or "").strip() + if not slug: + return ("unspecified", "unspecified") + if "/" in slug: + provider, _, name = slug.partition("/") + return (name or "unspecified", provider or "unspecified") + return (slug, "unspecified") + + +def read_number(value): + """A usable numeric value, or None. Booleans are numbers in Python and are not + usable ones here.""" + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + return value + + +def read_messages(path): + """Parse a file that is a JSON array, a single JSON object, or JSON-lines into a + list of dicts. Empty on anything unreadable. This mirrors review-diagnostics.py's + tolerant parsing: both tools' output formats have moved before.""" + if not path: + return [] + try: + with open(path) as handle: + text = handle.read() + except OSError: + return [] + try: + parsed = json.loads(text) + except json.JSONDecodeError: + parsed = [] + for line in text.splitlines(): + if not line.strip(): + continue + try: + parsed.append(json.loads(line)) + except json.JSONDecodeError: + continue + if isinstance(parsed, dict): + parsed = [parsed] + if not isinstance(parsed, list): + return [] + return [item for item in parsed if isinstance(item, dict)] + + +def openrouter_cost(usage): + """The amount OpenRouter reports charging for a call, from its usage object, or + None when the field didn't survive into what the tool recorded. `cost` is the + total billed; `cost_details.upstream_inference_cost` is the provider's share of + it and is the fallback when only the breakdown came through.""" + if not isinstance(usage, dict): + return None + cost = read_number(usage.get("cost")) + if cost is not None: + return cost + details = usage.get("cost_details") + if isinstance(details, dict): + return read_number(details.get("upstream_inference_cost")) + return None + + +def collect(usage, fields): + """Whichever of `fields` the usage object actually carries, under Datadog's + metric names.""" + counts = {} + if not isinstance(usage, dict): + return counts + for source, metric in fields.items(): + value = read_number(usage.get(source)) + if value is not None: + counts[metric] = int(value) + return counts + + +def claude_usage(path): + """Claude Code's execution file -> (token counts, reported cost or None). The + terminal "result" message carries `usage`; only that object is read.""" + results = [m for m in read_messages(path) if m.get("type") == "result"] + if not results: + return ({}, None) + usage = results[-1].get("usage") + return (collect(usage, CLAUDE_USAGE_FIELDS), openrouter_cost(usage)) + + +def find_rollout(directory): + """The most recently modified rollout file under Codex's session directory, which + nests them by date. None when the directory is absent or holds none.""" + newest = None + for root, _, names in os.walk(directory or ""): + for name in names: + if not (name.startswith("rollout-") and name.endswith(".jsonl")): + continue + path = os.path.join(root, name) + try: + stamp = os.path.getmtime(path) + except OSError: + continue + if newest is None or stamp > newest[0]: + newest = (stamp, path) + return newest[1] if newest else None + + +def codex_usage(directory): + """Codex's session rollout -> (token counts, reported cost or None). Codex logs a + running total after every turn under a `token_count` event; the last one is the + total for the session.""" + path = find_rollout(directory) + if not path: + return ({}, None) + totals = None + for event in read_messages(path): + payload = event.get("payload") + if not isinstance(payload, dict) or payload.get("type") != "token_count": + continue + info = payload.get("info") + if isinstance(info, dict) and isinstance(info.get("total_token_usage"), dict): + totals = info["total_token_usage"] + if totals is None: + return ({}, None) + return (collect(totals, CODEX_USAGE_FIELDS), openrouter_cost(totals)) + + +def build_span_fields(slug, counts, cost): + """Pure assembly of the span fields the workflow splices in: a catalog-matching + model identity, whichever token counts are known, and a reported cost when there + is one. `total_tokens` is Datadog's own metric name, so it is spelled out rather + than left for Datadog to infer.""" + model_name, provider = split_model(slug) + metrics = dict(counts) + if "input_tokens" in metrics and "output_tokens" in metrics: + metrics["total_tokens"] = metrics["input_tokens"] + metrics["output_tokens"] + if cost is not None: + metrics["total_cost"] = cost + return { + "model_name": model_name, + "model_provider": provider, + "gateway": GATEWAY, + "metrics": metrics, + } + + +def main(argv): + pass_name = argv[1] if len(argv) > 1 else "" + slug = argv[2] if len(argv) > 2 else "" + source = argv[3] if len(argv) > 3 else "" + if pass_name == "codex": + counts, cost = codex_usage(source or DEFAULT_ROLLOUT_DIR) + elif pass_name == "claude": + counts, cost = claude_usage(source) + else: + counts, cost = ({}, None) + print(json.dumps(build_span_fields(slug, counts, cost))) + return 0 + + +if __name__ == "__main__": + sys.exit(main(sys.argv)) diff --git a/tests/test_llm_usage.py b/tests/test_llm_usage.py new file mode 100644 index 0000000..9bd63e4 --- /dev/null +++ b/tests/test_llm_usage.py @@ -0,0 +1,180 @@ +import importlib.util +import json +import os +from pathlib import Path +import tempfile +import unittest + +SCRIPT = Path(__file__).resolve().parents[1] / 'scripts/llm-usage.py' +spec = importlib.util.spec_from_file_location('llm_usage', SCRIPT) +llm_usage = importlib.util.module_from_spec(spec) +spec.loader.exec_module(llm_usage) + + +def write(content, name='execution.json'): + path = Path(tempfile.mkdtemp()) / name + path.write_text(content) + return str(path) + + +def write_rollout(events, name='rollout-abc.jsonl', subdir='2026/09/10'): + directory = Path(tempfile.mkdtemp()) / subdir + directory.mkdir(parents=True) + (directory / name).write_text('\n'.join(json.dumps(e) for e in events)) + return str(Path(directory).parents[2]) + + +def token_count(**usage): + return {'type': 'event_msg', + 'payload': {'type': 'token_count', 'info': {'total_token_usage': usage}}} + + +class SplitModelTest(unittest.TestCase): + def test_splits_openrouter_slug_into_catalog_name_and_provider(self): + self.assertEqual(llm_usage.split_model('openai/gpt-5.6-sol'), + ('gpt-5.6-sol', 'openai')) + + def test_bare_model_reports_unspecified_provider(self): + self.assertEqual(llm_usage.split_model('gpt-5.6-sol'), + ('gpt-5.6-sol', 'unspecified')) + + def test_empty_slug_is_fully_unspecified(self): + self.assertEqual(llm_usage.split_model(''), ('unspecified', 'unspecified')) + self.assertEqual(llm_usage.split_model(None), ('unspecified', 'unspecified')) + + +class OpenrouterCostTest(unittest.TestCase): + def test_reads_the_charged_amount(self): + self.assertEqual(llm_usage.openrouter_cost({'cost': 0.0142}), 0.0142) + + def test_falls_back_to_the_upstream_share_when_only_the_breakdown_survives(self): + usage = {'cost_details': {'upstream_inference_cost': 0.009}} + self.assertEqual(llm_usage.openrouter_cost(usage), 0.009) + + def test_prefers_the_total_charge_over_the_upstream_share(self): + usage = {'cost': 0.0142, 'cost_details': {'upstream_inference_cost': 0.009}} + self.assertEqual(llm_usage.openrouter_cost(usage), 0.0142) + + def test_absent_cost_is_none_not_zero(self): + self.assertIsNone(llm_usage.openrouter_cost({'input_tokens': 10})) + self.assertIsNone(llm_usage.openrouter_cost(None)) + + def test_rejects_a_non_numeric_cost(self): + self.assertIsNone(llm_usage.openrouter_cost({'cost': 'free'})) + self.assertIsNone(llm_usage.openrouter_cost({'cost': True})) + + +class ClaudeUsageTest(unittest.TestCase): + def test_reads_counts_and_reported_cost_from_the_result_message(self): + path = write(json.dumps([ + {'type': 'assistant'}, + {'type': 'result', 'usage': {'input_tokens': 1000, 'output_tokens': 200, + 'cache_read_input_tokens': 800, + 'cache_creation_input_tokens': 50, + 'cost': 0.0031}}, + ])) + counts, cost = llm_usage.claude_usage(path) + self.assertEqual(counts, {'input_tokens': 1000, 'output_tokens': 200, + 'cache_read_input_tokens': 800, + 'cache_write_input_tokens': 50}) + self.assertEqual(cost, 0.0031) + + def test_ignores_claude_codes_own_cost_estimate(self): + path = write(json.dumps([{'type': 'result', 'total_cost_usd': 9.99, + 'usage': {'input_tokens': 5}}])) + _, cost = llm_usage.claude_usage(path) + self.assertIsNone(cost) + + def test_takes_the_last_result_message(self): + path = write('\n'.join([ + json.dumps({'type': 'result', 'usage': {'input_tokens': 1}}), + json.dumps({'type': 'result', 'usage': {'input_tokens': 7}}), + ])) + counts, _ = llm_usage.claude_usage(path) + self.assertEqual(counts['input_tokens'], 7) + + def test_missing_file_yields_nothing(self): + self.assertEqual(llm_usage.claude_usage('/nonexistent/execution.json'), ({}, None)) + self.assertEqual(llm_usage.claude_usage(''), ({}, None)) + + def test_unparseable_content_yields_nothing(self): + self.assertEqual(llm_usage.claude_usage(write('not json at all')), ({}, None)) + + def test_skips_non_numeric_counts(self): + path = write(json.dumps({'type': 'result', 'usage': { + 'input_tokens': 'lots', 'output_tokens': True, + 'cache_read_input_tokens': 5}})) + counts, _ = llm_usage.claude_usage(path) + self.assertEqual(counts, {'cache_read_input_tokens': 5}) + + +class CodexUsageTest(unittest.TestCase): + def test_reads_the_running_totals_from_the_rollout(self): + directory = write_rollout([ + token_count(input_tokens=100, output_tokens=10), + token_count(input_tokens=31751, cached_input_tokens=14720, + output_tokens=2367, total_tokens=34118), + ]) + counts, cost = llm_usage.codex_usage(directory) + self.assertEqual(counts, {'input_tokens': 31751, 'output_tokens': 2367, + 'cache_read_input_tokens': 14720}) + self.assertIsNone(cost) + + def test_takes_the_most_recent_rollout_when_several_exist(self): + directory = write_rollout([token_count(input_tokens=1)]) + newer = Path(directory) / '2026/09/11' + newer.mkdir(parents=True) + path = newer / 'rollout-def.jsonl' + path.write_text(json.dumps(token_count(input_tokens=42))) + os.utime(path, (2_000_000_000, 2_000_000_000)) + counts, _ = llm_usage.codex_usage(directory) + self.assertEqual(counts['input_tokens'], 42) + + def test_rollout_without_token_counts_yields_nothing(self): + directory = write_rollout([{'type': 'event_msg', 'payload': {'type': 'agent_message'}}]) + self.assertEqual(llm_usage.codex_usage(directory), ({}, None)) + + def test_missing_session_directory_yields_nothing(self): + self.assertEqual(llm_usage.codex_usage('/nonexistent/sessions'), ({}, None)) + self.assertEqual(llm_usage.codex_usage(''), ({}, None)) + + +class BuildSpanFieldsTest(unittest.TestCase): + def test_reports_catalog_identity_and_derived_total(self): + fields = llm_usage.build_span_fields( + 'openai/gpt-5.6-sol', {'input_tokens': 1000, 'output_tokens': 200}, None) + self.assertEqual(fields['model_name'], 'gpt-5.6-sol') + self.assertEqual(fields['model_provider'], 'openai') + self.assertEqual(fields['gateway'], 'openrouter') + self.assertEqual(fields['metrics']['total_tokens'], 1200) + + def test_omits_cost_when_none_was_reported(self): + fields = llm_usage.build_span_fields( + 'openai/gpt-5.6-sol', {'input_tokens': 1000, 'output_tokens': 200}, None) + self.assertNotIn('total_cost', fields['metrics']) + + def test_reports_a_reported_cost_under_datadogs_metric_name(self): + fields = llm_usage.build_span_fields( + 'deepseek/deepseek-v4.1-flash', {'input_tokens': 1000}, 0.0031) + self.assertEqual(fields['metrics']['total_cost'], 0.0031) + + def test_omits_the_derived_total_when_a_count_is_missing(self): + fields = llm_usage.build_span_fields('x/y', {'input_tokens': 1000}, None) + self.assertNotIn('total_tokens', fields['metrics']) + + def test_no_usage_at_all_still_yields_a_usable_span(self): + fields = llm_usage.build_span_fields('openai/gpt-5.6-sol', {}, None) + self.assertEqual(fields['metrics'], {}) + self.assertEqual(fields['model_name'], 'gpt-5.6-sol') + + +class MainTest(unittest.TestCase): + def test_unknown_pass_name_emits_an_empty_but_valid_object(self): + llm_usage.main(['llm-usage.py', 'nonsense', 'openai/gpt-5.6-sol']) + + def test_exits_zero_with_no_arguments(self): + self.assertEqual(llm_usage.main(['llm-usage.py']), 0) + + +if __name__ == '__main__': + unittest.main()