ZacharyZhang-NY/omarchycn
3
4
Fork 0
Code Issues Pull Requests 1 Actions Packages Projects Releases 5 Wiki Activity

Fix model usage token counting

Browse Source
This commit is contained in:
David Heinemeier Hansson
2026-07-29 18:48:15 -07:00
parent f8835df644
commit 4d05d39d33
3 changed files with 92 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files
shell/plugins/model-usage/scripts/codex_usage_scanner.py
+10 -5
View File
@@ -213,12 +213,17 @@ def scan_native_codex_sessions():
if payload.get("type") != "token_count":
continue
info = payload.get("info") or {}
usage = info.get("total_token_usage") or {}
input_tokens = number(usage.get("input_tokens"))
output_tokens = number(usage.get("output_tokens")) + number(usage.get("reasoning_output_tokens"))
# total_token_usage is cumulative for the session. Adding every
# snapshot makes usage grow quadratically, so count the last turn.
usage = info.get("last_token_usage") or {}
cache_read = number(usage.get("cached_input_tokens"))
cache_write = 0
if not (input_tokens or output_tokens or cache_read):
cache_write = number(usage.get("cache_write_input_tokens"))
# Cached tokens are included in input_tokens, and reasoning tokens
# are included in output_tokens. Keep the cache split without
# counting either category twice.
input_tokens = max(0, number(usage.get("input_tokens")) - cache_read - cache_write)
output_tokens = number(usage.get("output_tokens"))
if not (input_tokens or output_tokens or cache_read or cache_write):
continue
day = local_day(entry.get("timestamp") or path.stat().st_mtime)
add_usage(day, str(path), current_model, input_tokens, output_tokens, cache_read, cache_write)
test/shell.d/model-usage-claude-scanner-test.sh
+30
View File
@@ -0,0 +1,30 @@
#!/bin/bash
source "$(dirname "$0")/base-test.sh"
require_command jq
require_command python3
TEST_HOME=$(mktemp -d)
trap 'rm -rf "$TEST_HOME"' EXIT
projects="$TEST_HOME/.claude/projects/example"
mkdir -p "$projects"
timestamp="$(date +%Y-%m-%d)T12:00:00Z"
cat >"$projects/session.jsonl" <<EOF
{"timestamp":"$timestamp","type":"assistant","sessionId":"session-1","uuid":"event-1","message":{"id":"message-1","role":"assistant","model":"claude-test","usage":{"input_tokens":2,"cache_creation_input_tokens":28857,"cache_read_input_tokens":0,"output_tokens":231}}}
{"timestamp":"$timestamp","type":"assistant","sessionId":"session-1","uuid":"event-2","message":{"id":"message-1","role":"assistant","model":"claude-test","usage":{"input_tokens":2,"cache_creation_input_tokens":28857,"cache_read_input_tokens":0,"output_tokens":231}}}
{"timestamp":"$timestamp","type":"assistant","sessionId":"session-1","uuid":"event-3","message":{"id":"message-2","role":"assistant","model":"claude-test","usage":{"input_tokens":2,"cache_creation_input_tokens":454,"cache_read_input_tokens":28857,"output_tokens":390}}}
EOF
result=$(HOME="$TEST_HOME" XDG_CACHE_HOME="$TEST_HOME/.cache" \
python3 "$ROOT/shell/plugins/model-usage/scripts/claude_usage_scanner.py" "$TEST_HOME/.claude/projects" --force)
[[ $(jq -r '.todayTotalTokens' <<<"$result") == "58793" ]] ||
fail "Claude scanner counts each API message once" "$result"
pass "Claude scanner counts each API message once"
[[ $(jq -c '.modelUsage["claude-test"]' <<<"$result") == '{"cacheCreationInputTokens":29311,"cacheReadInputTokens":28857,"inputTokens":4,"outputTokens":621}' ]] ||
fail "Claude scanner keeps mutually exclusive token categories" "$result"
pass "Claude scanner keeps mutually exclusive token categories"
test/shell.d/model-usage-codex-scanner-test.sh
+52
View File
@@ -0,0 +1,52 @@
#!/bin/bash
source "$(dirname "$0")/base-test.sh"
require_command jq
require_command python3
TEST_HOME=$(mktemp -d)
trap 'rm -rf "$TEST_HOME"' EXIT
mkdir -p "$TEST_HOME/.codex/sessions/$(date +%Y/%m/%d)" "$TEST_HOME/bin"
cat >"$TEST_HOME/bin/codex" <<'EOF'
#!/bin/bash
while read -r request; do
id=$(jq -r '.id // empty' <<<"$request")
method=$(jq -r '.method // empty' <<<"$request")
case "$method" in
initialize)
jq -cn --argjson id "$id" '{id: $id, result: {}}'
;;
account/read)
jq -cn --argjson id "$id" '{id: $id, result: {account: {}}}'
;;
account/rateLimits/read)
jq -cn --argjson id "$id" '{id: $id, result: {rateLimits: {}}}'
;;
esac
done
EOF
chmod +x "$TEST_HOME/bin/codex"
timestamp="$(date +%Y-%m-%d)T12:00:00Z"
session="$TEST_HOME/.codex/sessions/$(date +%Y/%m/%d)/rollout.jsonl"
cat >"$session" <<EOF
{"timestamp":"$timestamp","type":"turn_context","payload":{"model":"gpt-test"}}
{"timestamp":"$timestamp","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":100,"cached_input_tokens":60,"output_tokens":20,"reasoning_output_tokens":5,"total_tokens":120},"last_token_usage":{"input_tokens":100,"cached_input_tokens":60,"output_tokens":20,"reasoning_output_tokens":5,"total_tokens":120}}}}
{"timestamp":"$timestamp","type":"event_msg","payload":{"type":"token_count","info":{"total_token_usage":{"input_tokens":180,"cached_input_tokens":110,"output_tokens":30,"reasoning_output_tokens":8,"total_tokens":210},"last_token_usage":{"input_tokens":80,"cached_input_tokens":50,"output_tokens":10,"reasoning_output_tokens":3,"total_tokens":90}}}}
EOF
result=$(HOME="$TEST_HOME" CODEX_HOME="$TEST_HOME/.codex" PATH="$TEST_HOME/bin:$PATH" \
python3 "$ROOT/shell/plugins/model-usage/scripts/codex_usage_scanner.py")
[[ $(jq -r '.todayTotalTokens' <<<"$result") == "210" ]] ||
fail "Codex scanner counts each turn once" "$result"
pass "Codex scanner counts each turn once"
[[ $(jq -c '.modelUsage["gpt-test"]' <<<"$result") == '{"inputTokens":70,"outputTokens":30,"cacheReadInputTokens":110,"cacheCreationInputTokens":0}' ]] ||
fail "Codex scanner does not double-count cache or reasoning tokens" "$result"
pass "Codex scanner does not double-count cache or reasoning tokens"