Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 25 additions & 7 deletions src/claude_statusbar/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -727,6 +727,22 @@ def _get(key):
return None


def _exact_used_tokens(stdin_data: Dict[str, Any]) -> Optional[int]:
"""Context tokens as Claude Code reported them, or None when unavailable.

None means "no usable signal" — an older Claude Code that omits the
totals, a relay payload that zeroes them, or a malformed value. Callers
fall back to deriving the count from the percentage.
"""
total = 0
for key in ('total_input_tokens', 'total_output_tokens'):
try:
total += int(stdin_data.get(key, 0) or 0)
except (TypeError, ValueError):
return None
return total or None
Comment on lines +730 to +743

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win

Require both complete, non-negative token totals.

.get(key, 0) or 0 silently treats a missing or empty field as zero, so a payload with only one total is used as an exact count instead of falling back to the percentage. Negative values and fractional floats are also accepted by int(). Require both fields and reject non-integer or negative values before summing; add a regression test for a single missing total.

As per path instructions, malformed Claude JSONL must not produce incorrect token calculations.

Proposed fix
-    total = 0
+    totals = []
     for key in ('total_input_tokens', 'total_output_tokens'):
+        if key not in stdin_data:
+            return None
+        raw = stdin_data[key]
+        if raw is None or isinstance(raw, bool):
+            return None
         try:
-            total += int(stdin_data.get(key, 0) or 0)
+            value = int(raw)
         except (TypeError, ValueError):
             return None
-    return total or None
+        if value < 0 or (isinstance(raw, float) and not raw.is_integer()):
+            return None
+        totals.append(value)
+    total = sum(totals)
+    return total or None
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
def _exact_used_tokens(stdin_data: Dict[str, Any]) -> Optional[int]:
"""Context tokens as Claude Code reported them, or None when unavailable.
None means "no usable signal"an older Claude Code that omits the
totals, a relay payload that zeroes them, or a malformed value. Callers
fall back to deriving the count from the percentage.
"""
total = 0
for key in ('total_input_tokens', 'total_output_tokens'):
try:
total += int(stdin_data.get(key, 0) or 0)
except (TypeError, ValueError):
return None
return total or None
def _exact_used_tokens(stdin_data: Dict[str, Any]) -> Optional[int]:
"""Context tokens as Claude Code reported them, or None when unavailable.
None means "no usable signal"an older Claude Code that omits the
totals, a relay payload that zeroes them, or a malformed value. Callers
fall back to deriving the count from the percentage.
"""
totals = []
for key in ('total_input_tokens', 'total_output_tokens'):
if key not in stdin_data:
return None
raw = stdin_data[key]
if raw is None or isinstance(raw, bool):
return None
try:
value = int(raw)
except (TypeError, ValueError):
return None
if value < 0 or (isinstance(raw, float) and not raw.is_integer()):
return None
totals.append(value)
total = sum(totals)
return total or None
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@src/claude_statusbar/core.py` around lines 730 - 743, Update
_exact_used_tokens to require both total_input_tokens and total_output_tokens to
be present, valid non-negative integers before summing; reject missing, empty,
malformed, negative, and fractional values by returning None so callers use the
percentage fallback. Add a regression test covering a payload with one missing
total.

Source: Path instructions



def _context_window_usage(stdin_data: Dict[str, Any],
env=None) -> Tuple[Optional[float], int, int]:
"""Return (ctx_pct, ctx_size, ctx_used) for renderer/model suffix.
Expand All @@ -752,13 +768,15 @@ def _context_window_usage(stdin_data: Dict[str, Any],
except (TypeError, ValueError):
ctx_pct = None

if ctx_pct is not None:
ctx_used = int(ctx_size_f * ctx_pct / 100)
else:
ctx_used = (
stdin_data.get('total_input_tokens', 0)
+ stdin_data.get('total_output_tokens', 0)
)
# Exact totals beat the integer percentage. Claude Code reports
# used_percentage as a whole number, so deriving tokens from it quantises
# the readout to 1% of the window — 2k on a 200k model (invisible), but
# 10k on a 1M-context one, where the bar steps 60.0k → 70.0k and can sit
# ~10k off the truth. `ctx_pct` itself stays as reported so the percentage
# and its severity colour keep matching what Claude Code shows.
ctx_used = _exact_used_tokens(stdin_data)
if ctx_used is None:
ctx_used = int(ctx_size_f * ctx_pct / 100) if ctx_pct is not None else 0

# Env override (#29): used tokens stay what stdin reported; the window and
# the percentage are re-derived against the real (env-forced) size.
Expand Down
31 changes: 31 additions & 0 deletions tests/test_core_ctx_pct.py
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,37 @@ def test_normal_context_returns_float():
assert isinstance(out, float)


def test_used_tokens_come_from_exact_totals_not_rounded_pct():
"""used_percentage is a whole number; on a 1M window that quantises the
token readout to 10k steps. The exact totals must win."""
ctx_pct, _, ctx_used = core._context_window_usage({
"context_window_size": 1_000_000,
"context_used_pct": 6,
"total_input_tokens": 63_824,
"total_output_tokens": 0,
})
assert ctx_used == 63_824 # not 60_000
assert ctx_pct == 6.0 # percentage stays as Claude Code reported it


def test_used_tokens_fall_back_to_pct_when_totals_absent():
"""Older Claude Code / relay payloads omit the totals entirely."""
_, _, ctx_used = core._context_window_usage({
"context_window_size": 200_000,
"context_used_pct": 25,
})
assert ctx_used == 50_000


def test_malformed_totals_fall_back_to_pct():
_, _, ctx_used = core._context_window_usage({
"context_window_size": 200_000,
"context_used_pct": 25,
"total_input_tokens": "lots",
})
assert ctx_used == 50_000


def test_null_context_pct_is_unknown_not_error():
ctx_pct, ctx_size, ctx_used = core._context_window_usage({
"context_window_size": 1_000_000,
Expand Down
Loading