diff --git a/navi/core/compressor.py b/navi/core/compressor.py index 95c35cf..f5f08db 100644 --- a/navi/core/compressor.py +++ b/navi/core/compressor.py @@ -26,6 +26,7 @@ from navi.config import settings from .events import ContextCompressed from .token_calibration import calibration_for +from .context_builder import truncate_oversized if TYPE_CHECKING: from navi.profiles.base import AgentProfile @@ -1026,7 +1027,13 @@ ) session.context = new_context - session.context_token_count = self.estimate_tokens(new_context, model) + # The size the model is actually sent. A megabyte-long tool result stays + # in stored history (the user still sees it) but build() truncates it per + # message before the call — counting the raw message reported 817k tokens + # against a 128k window the model never received. + session.context_token_count = self.estimate_tokens( + truncate_oversized(new_context), model + ) # The context just shrank, so the recorded real-token baseline is stale. # Drop it — the next LLM call re-establishes a fresh one. self._real_baseline = None diff --git a/navi/core/context_builder.py b/navi/core/context_builder.py index f423676..ad6cb0e 100644 --- a/navi/core/context_builder.py +++ b/navi/core/context_builder.py @@ -66,6 +66,35 @@ return content[:half] + marker + content[-half:] +def truncate_oversized(conv: list[Message]) -> list[Message]: + """Head/tail-truncate oversized tool/assistant messages for the LLM view. + + A single huge tool result (e.g. reading a large file) can alone exceed + the context window, and the turn/message-count-based compressor cannot + shrink a single message. This per-build pass caps each truncatable + message at ``context_message_token_budget`` tokens (0 = num_ctx // 6), + keeping head + tail so the model still sees the start and end of the + result. Copies are returned — stored history (session.messages) is never + mutated, so the user still sees full content and reloads are unaffected. + + Module-level because the compressor reports the session's context size in + the same units: a raw context holding a truncated-away megabyte is not what + the model is sent, and counting it as such shows a scary 800k over a 128k + window. + """ + budget = _config.settings.context_message_token_budget + if budget <= 0: + budget = max(1, _config.settings.ollama_num_ctx // 6) + cap_chars = budget * 3 + out: list[Message] = [] + for m in conv: + if m.role in _TRUNCATABLE_ROLES and len(m.content or "") > cap_chars: + out.append(m.model_copy(update={"content": _truncate_content(m.content or "", budget)})) + else: + out.append(m) + return out + + def _fact_value_is_abs_path_outside(value: str, cwd: str) -> bool: """True when a fact's value is an absolute filesystem path outside `cwd`. @@ -449,27 +478,8 @@ return Message(role="system", content="\n".join(lines)) def _truncate_oversized(self, conv: list[Message]) -> list[Message]: - """Head/tail-truncate oversized tool/assistant messages for the LLM view. - - A single huge tool result (e.g. reading a large file) can alone exceed - the context window, and the turn/message-count-based compressor cannot - shrink a single message. This per-build pass caps each truncatable - message at ``context_message_token_budget`` tokens (0 = num_ctx // 6), - keeping head + tail so the model still sees the start and end of the - result. Copies are returned — stored history (session.messages) is never - mutated, so the user still sees full content and reloads are unaffected. - """ - budget = _config.settings.context_message_token_budget - if budget <= 0: - budget = max(1, _config.settings.ollama_num_ctx // 6) - cap_chars = budget * 3 - out: list[Message] = [] - for m in conv: - if m.role in _TRUNCATABLE_ROLES and len(m.content or "") > cap_chars: - out.append(m.model_copy(update={"content": _truncate_content(m.content or "", budget)})) - else: - out.append(m) - return out + """Per-build pass; see ``truncate_oversized`` for the rule.""" + return truncate_oversized(conv) def build( self, diff --git a/tests/unit/core/test_compressor.py b/tests/unit/core/test_compressor.py index 0914bdc..355052a 100644 --- a/tests/unit/core/test_compressor.py +++ b/tests/unit/core/test_compressor.py @@ -1354,3 +1354,77 @@ ) assert event is None assert store.context_flips == [] + + +class TestCompressedTokenCountIsTheEffectiveView: + """The counter the UI shows is the prompt the model is sent, not the raw + stored context. A 1.47MB code_exec listing once made a prod session report + 817 288 / 131 072 (100%) right after a compression pass, while build() + truncates that message per call — the model never saw 800k tokens.""" + + @staticmethod + def _override_settings(): + import navi.config as _config + import navi.core.compressor as _cm + from navi.config import Settings + + new = Settings( + _env_file=None, + ollama_num_ctx=10_000, + context_message_token_budget=100, # ~300 chars per message + context_compression_target=0.65, + ) + prev_config, prev_cm = _config.settings, _cm.settings + _config.settings = new + _cm.settings = new + return prev_config, prev_cm + + @staticmethod + def _restore(prev_config, prev_cm): + import navi.config as _config + import navi.core.compressor as _cm + + _config.settings = prev_config + _cm.settings = prev_cm + + @pytest.mark.asyncio + async def test_oversized_kept_message_is_counted_as_the_model_sees_it(self): + from navi.core.context_builder import truncate_oversized + from navi.core.session import Session + + prev_config, prev_cm = self._override_settings() + try: + store = _RecordingStore() + c = ContextCompressor() + session = Session(profile_id="test") + session.context = [Message(role="system", content="sys")] + for i in range(10): + session.context.append(Message(role="user", content=f"u{i}")) + session.context.append(Message(role="assistant", content="a" * 400)) + big = Message(role="tool", name="code_exec", content="x" * 1_000_000) + session.context.append(big) # most recent turn — kept verbatim + session.messages = list(session.context) + + event = await c.compress_and_save_session( + session=session, + session_store=store, + llm=FakeLLMBackend(responses=["summary"]), + model="m", + temperature=0.3, + session_id=session.id, + reason="pytest", + keep_recent=1, + window=10_000, + ) + + assert event is not None + # Stored history keeps the whole result: the user still sees it. + assert big in session.context + raw = c.estimate_tokens(session.context, "m") + effective = c.estimate_tokens(truncate_oversized(session.context), "m") + assert session.context_token_count == effective + assert session.context_token_count == event.context_tokens + # The raw estimate is the scary number; the stored one is not it. + assert raw > effective * 10 + finally: + self._restore(prev_config, prev_cm)