From 12270ca9fab67f17617be4fbc5ead975bdcb7e6c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 22 May 2026 15:48:28 +0000 Subject: [PATCH] Round estimate_tokens up so over-budget prompts don't skip compaction MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit `estimate_tokens` used floor division (`// 4`) on the char count, which systematically underestimates any non-multiple-of-4 message length. A 5-char message returned 1 token instead of 2, so a prompt sitting just over the budget could appear under threshold and bypass compaction — the exact failure the module exists to prevent. Switch to ceil-style rounding so the heuristic stays conservative. Adds regression tests covering the off-by-one and the single-char case (where floor returned 0 tokens for non-empty content). --- studio/backend/core/inference/context_compaction.py | 7 ++++++- studio/backend/tests/test_context_compaction.py | 11 +++++++++++ 2 files changed, 17 insertions(+), 1 deletion(-) diff --git a/studio/backend/core/inference/context_compaction.py b/studio/backend/core/inference/context_compaction.py index 8428d0cd53..18c36bb07c 100644 --- a/studio/backend/core/inference/context_compaction.py +++ b/studio/backend/core/inference/context_compaction.py @@ -65,7 +65,12 @@ def estimate_tokens(messages: list[dict]) -> int: args = (tc.get("function") or {}).get("arguments") if isinstance(args, str): total_chars += len(args) - return total_chars // _CHARS_PER_TOKEN + # Ceil-style division: floor (`// 4`) would systematically + # underestimate non-multiple-of-4 lengths, letting just-over-budget + # prompts appear under threshold and bypass compaction — the exact + # failure this module exists to prevent. Round up so the heuristic + # stays on the conservative side described above. + return -(-total_chars // _CHARS_PER_TOKEN) def _is_multimodal(msg: dict) -> bool: diff --git a/studio/backend/tests/test_context_compaction.py b/studio/backend/tests/test_context_compaction.py index 03e688e78a..c42b39d9ad 100644 --- a/studio/backend/tests/test_context_compaction.py +++ b/studio/backend/tests/test_context_compaction.py @@ -47,6 +47,17 @@ class TestEstimateTokens: msgs = [_msg("user", "abcd")] assert estimate_tokens(msgs) == 1 + def test_non_multiple_of_4_rounds_up(self): + # Regression: floor (`// 4`) would estimate 1 for "abcde" and + # let an over-budget message slip past the threshold check. + # Ceil keeps the heuristic conservative. + msgs = [_msg("user", "abcde")] # 5 chars + assert estimate_tokens(msgs) == 2 + + def test_single_char_rounds_up_to_one_token(self): + msgs = [_msg("user", "a")] + assert estimate_tokens(msgs) == 1 + def test_multimodal_text_part_counts(self): msgs = [ {