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 = [ {