diff --git a/bin/omarchy-agent-usage-claude b/bin/omarchy-agent-usage-claude index 5b3634aef1c..012cf9c4dac 100755 --- a/bin/omarchy-agent-usage-claude +++ b/bin/omarchy-agent-usage-claude @@ -123,7 +123,15 @@ def scan_projects(projects_path: Path) -> dict[str, Any]: recent_dates = recent_date_strings() recent = {day: {"date": day, "messageCount": 0} for day in recent_dates} - seen: set[str] = set() + # A streamed response is several lines sharing one message id, and the + # first line's output_tokens is a placeholder while the last carries the + # real count. Keep each message's line with the highest output count + # (the last one scanned wins a tie) and tally once all lines are in. The + # whole line is kept: a response that fell back to another model + # mid-stream has lines that are separate snapshots, and mixing fields + # across them over-counts cache tokens and credits the wrong model. + # Tuple: (output_tokens, model, day, session_key, input, cache_read, cache_write). + seen: dict[str, tuple[int, str, str, str, int, int, int]] = {} sessions: set[str] = set() active_days: set[str] = set() today_sessions: set[str] = set() @@ -158,44 +166,48 @@ def scan_projects(projects_path: Path) -> dict[str, Any]: message_id = message.get("id") or entry.get("messageId") or "" unique_key = str(message_id) if message_id else f"{path}:{entry.get('uuid') or entry.get('requestId') or line_number}" - if unique_key in seen: - continue - seen.add(unique_key) input_tokens = usage_token(usage, "input_tokens", "inputTokens") output_tokens = usage_token(usage, "output_tokens", "outputTokens") cache_read = usage_token(usage, "cache_read_input_tokens", "cacheReadInputTokens") cache_write = usage_token(usage, "cache_creation_input_tokens", "cacheCreationInputTokens") - total = input_tokens + output_tokens + cache_read + cache_write - if total <= 0: + if input_tokens + output_tokens + cache_read + cache_write <= 0: + continue + + held = seen.get(unique_key) + if held is not None and output_tokens < held[0]: continue model = str(message.get("model") or entry.get("model") or "claude") day = local_date_from_timestamp(entry.get("timestamp") or message.get("timestamp")) session_key = str(entry.get("sessionId") or path) - sessions.add(session_key) - active_days.add(day) - prompts += 1 - - bucket = usage_by_model.setdefault(model, empty_bucket()) - bucket["inputTokens"] += input_tokens - bucket["outputTokens"] += output_tokens - bucket["cacheReadInputTokens"] += cache_read - bucket["cacheCreationInputTokens"] += cache_write - - if day in recent: - # recentDays.messageCount is actually a token total, despite the - # legacy name shared with synced snapshots. - recent[day]["messageCount"] += total - - if day == today: - today_prompt_count += 1 - today_sessions.add(session_key) - today_token_total += total - today_tokens[model] = today_tokens.get(model, 0) + total + seen[unique_key] = (output_tokens, model, day, session_key, input_tokens, cache_read, cache_write) except Exception as exc: print(f"Ignoring unreadable Claude project file {path}: {exc}", file=sys.stderr) + for output_tokens, model, day, session_key, input_tokens, cache_read, cache_write in seen.values(): + total = input_tokens + output_tokens + cache_read + cache_write + sessions.add(session_key) + active_days.add(day) + prompts += 1 + + bucket = usage_by_model.setdefault(model, empty_bucket()) + bucket["inputTokens"] += input_tokens + bucket["outputTokens"] += output_tokens + bucket["cacheReadInputTokens"] += cache_read + bucket["cacheCreationInputTokens"] += cache_write + + if day in recent: + # recentDays.messageCount is actually a token total, despite the + # legacy name shared with synced snapshots. + recent[day]["messageCount"] += total + + if day == today: + today_prompt_count += 1 + today_sessions.add(session_key) + today_token_total += total + today_tokens[model] = today_tokens.get(model, 0) + total + return { "todayPrompts": today_prompt_count, "todaySessions": len(today_sessions), diff --git a/test/shell.d/agent-usage-claude-scanner-test.sh b/test/shell.d/agent-usage-claude-scanner-test.sh index e1407f3cf2e..2f6c9c83ab2 100644 --- a/test/shell.d/agent-usage-claude-scanner-test.sh +++ b/test/shell.d/agent-usage-claude-scanner-test.sh @@ -29,14 +29,46 @@ pass "Claude collector counts each API message once" fail "Claude collector keeps mutually exclusive token categories" "$result" pass "Claude collector keeps mutually exclusive token categories" + [[ $(jq -r '.id + "/" + .usageStatusText' <<<"$result") == "claude/Waiting for auth" ]] || fail "Claude collector identifies itself and reports missing auth" "$result" pass "Claude collector identifies itself and reports missing auth" +# A streamed response is several lines sharing one message id. The first +# line's output_tokens is a placeholder and the last line has the real count, +# so the message is counted from the line with the highest output. +STREAM_HOME=$(mktemp -d) +trap 'rm -rf "$TEST_HOME" "$STREAM_HOME"' EXIT +stream_projects="$STREAM_HOME/.claude/projects/example" +mkdir -p "$stream_projects" +cat >"$stream_projects/session.jsonl" <"$PI_HOME/.pi/agent/sessions/project/pi.jsonl" <