diff --git a/src/claude_statusbar/core.py b/src/claude_statusbar/core.py index 5dd75b8..e9fe66d 100755 --- a/src/claude_statusbar/core.py +++ b/src/claude_statusbar/core.py @@ -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 + + 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. @@ -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. diff --git a/tests/test_core_ctx_pct.py b/tests/test_core_ctx_pct.py index 8da4a02..592da77 100644 --- a/tests/test_core_ctx_pct.py +++ b/tests/test_core_ctx_pct.py @@ -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,