diff --git a/navi/core/agent.py b/navi/core/agent.py index b51ca07..a680ccc 100644 --- a/navi/core/agent.py +++ b/navi/core/agent.py @@ -237,7 +237,11 @@ raise NothingToCompactError("Context compression is disabled.") yield CompressionStarted( - context_tokens=self._compressor.estimate_context_tokens(session.context), + # Real-baseline estimate when available (see _compression_events_preturn); + # heuristic is only the fallback before the first LLM call of a session. + context_tokens=self._compressor.real_baseline_estimate( + session.context, session.context + ), max_context_tokens=settings.ollama_num_ctx, ) event = await self._compressor.compress_and_save_session( @@ -736,11 +740,19 @@ return TodoUpdated(session_id=session.id, tasks=tasks) async def _compression_events_preturn(self, session, llm, profile, session_id): + # real_baseline_estimate, not the chars//3 heuristic: the baseline is + # valid here (the context only grew since the last LLM call), and the + # heuristic undercounts code-heavy tool output — firing compression + # too late is exactly what the baseline machinery was built to fix. + # The midturn gate below does the same. + estimated_tokens = self._compressor.real_baseline_estimate( + session.context, session.context + ) if ( settings.context_compression_enabled and len(session.context) > 2 and should_compress( - self._compressor.estimate_context_tokens(session.context), + estimated_tokens, settings.ollama_num_ctx, settings.context_compression_threshold, ) @@ -759,7 +771,7 @@ ): return yield CompressionStarted( - context_tokens=self._compressor.estimate_context_tokens(session.context), + context_tokens=estimated_tokens, max_context_tokens=settings.ollama_num_ctx, ) event = await self._compressor.compress_and_save_session( diff --git a/navi/core/compressor.py b/navi/core/compressor.py index 47d7868..1b4dee5 100644 --- a/navi/core/compressor.py +++ b/navi/core/compressor.py @@ -16,6 +16,7 @@ import json import re +import structlog from datetime import datetime, timezone from pathlib import Path from typing import TYPE_CHECKING @@ -27,6 +28,8 @@ if TYPE_CHECKING: from navi.profiles.base import AgentProfile +log = structlog.get_logger() + _SUMMARY_SECTIONS = [ ("Goal", "One clear sentence describing what the user is trying to accomplish in this session. Include deadlines or constraints if stated."), @@ -272,6 +275,14 @@ ) if critical and len(result) <= 4000: preview = result + elif critical: + # Over the critical budget: keep head + tail halves rather + # than collapsing to the non-critical 800-char preview — + # the all-or-nothing cliff used to throw away 98% of a + # large file read (the common navi_code shape) while a + # 3999-char result survived whole. + half = 2000 + preview = result[:half] + "\n…[middle elided]…\n" + result[-half:] else: # Give the summarizer enough of the result to preserve # diagnostic detail (the exact error line, a grep match, @@ -302,7 +313,7 @@ # When existing summaries in to_summarize exceed this many chars combined, # run a quick meta-summary to consolidate them before the main compression pass. -_META_SUMMARY_THRESHOLD = _MAX_SUMMARY_INPUT_CHARS // 3 # 8_000 +_META_SUMMARY_THRESHOLD = _MAX_SUMMARY_INPUT_CHARS // 3 # 10_666 _META_SUMMARY_SYSTEM = ( "You are condensing multiple conversation summaries into a single compact summary. " @@ -361,6 +372,62 @@ return base + extra +def _plan_compression( + context: list[Message], + keep_recent: int, + keep_recent_messages: int | None = None, +) -> tuple[list[Message], list[Message], list[Message]] | None: + """Decide what to summarize and what to keep — the single decision point + shared by ``compress_context`` (real compression) and + ``ContextCompressor.would_compress`` (dry-run prediction), so the two can + never drift apart. + + Applies, in order: + 1. Turn-based partition (``partition_messages``). + 2. Target hysteresis (turn-based mode only): shrink the verbatim + keep-recent window until the kept region fits the target fraction of + the context window. A fixed keep_recent leaves roughly the same context + size after compression (just below the 90% trigger), so the next few + messages re-trigger it — the "125 -> 104, then 104 again after a couple + of messages" cycle. Dropping the kept size to the 65% target leaves + real headroom and folds the extra turns into the summary (same single + LLM call, no extra round-trip). Midturn mode (keep_recent_messages set) + already splits the in-flight turn aggressively, so a post-check in + compress_session truncates if it is still over target. + 3. Aggressive intra-turn fallback (midturn mode only): when the turn-based + partition found nothing, retry with keep_recent_messages=2. + + Returns (system_msgs, to_summarize, to_keep) or None when there is nothing + substantial to compress. + """ + system_msgs = [m for m in context if m.role == "system"] + to_summarize, to_keep = partition_messages( + context, keep_recent, keep_recent_messages=keep_recent_messages + ) + + if keep_recent_messages is None: + target_tokens = int(settings.ollama_num_ctx * settings.context_compression_target) + kept_tokens = ContextCompressor.estimate_context_tokens(system_msgs + to_keep) + _KEEP_RECENT_FLOOR = 2 + while kept_tokens > target_tokens and keep_recent > _KEEP_RECENT_FLOOR: + keep_recent -= 1 + to_summarize, to_keep = partition_messages( + context, keep_recent, keep_recent_messages=keep_recent_messages + ) + kept_tokens = ContextCompressor.estimate_context_tokens(system_msgs + to_keep) + + if len(to_summarize) < 2 and keep_recent_messages is not None and keep_recent_messages > 2: + to_summarize, to_keep = partition_messages( + context, + keep_recent, + keep_recent_messages=2, + ) + + if len(to_summarize) < 2: + return None + return system_msgs, to_summarize, to_keep + + async def compress_context( context: list[Message], llm: LLMBackend, @@ -393,48 +460,12 @@ effective_keep_recent = getattr(profile, "compression_keep_recent", None) or keep_recent effective_max_tokens = getattr(profile, "compression_max_tokens", None) or max_tokens - system_msgs = [m for m in context if m.role == "system"] - to_summarize, to_keep = partition_messages( - context, - effective_keep_recent, - keep_recent_messages=keep_recent_messages, + planned = _plan_compression( + context, effective_keep_recent, keep_recent_messages=keep_recent_messages ) - - # Target hysteresis (turn-based / preturn mode): shrink the verbatim - # keep-recent window until the kept region fits the target fraction of the - # context window. A fixed keep_recent leaves roughly the same context size - # after compression (just below the 90% trigger), so the next few messages - # re-trigger it — the "125 -> 104, then 104 again after a couple of - # messages" cycle. Dropping the kept size to the 65% target leaves real - # headroom and folds the extra turns into the summary (same single LLM - # call, no extra round-trip). Midturn mode (keep_recent_messages set) - # already splits the in-flight turn aggressively, so a post-check in - # compress_session truncates if it is still over target. - if keep_recent_messages is None: - target_tokens = int(settings.ollama_num_ctx * settings.context_compression_target) - kept_tokens = ContextCompressor.estimate_context_tokens(system_msgs + to_keep) - _KEEP_RECENT_FLOOR = 2 - while ( - kept_tokens > target_tokens and effective_keep_recent > _KEEP_RECENT_FLOOR - ): - effective_keep_recent -= 1 - to_summarize, to_keep = partition_messages( - context, effective_keep_recent, keep_recent_messages=keep_recent_messages - ) - kept_tokens = ContextCompressor.estimate_context_tokens(system_msgs + to_keep) - - # Fallback: if turn-based partition has nothing to compress but we are in - # mid-turn mode (keep_recent_messages set), try an aggressive intra-turn - # split keeping only the 2 newest messages of the current turn. - if len(to_summarize) < 2 and keep_recent_messages is not None and keep_recent_messages > 2: - to_summarize, to_keep = partition_messages( - context, - effective_keep_recent, - keep_recent_messages=2, - ) - - if len(to_summarize) < 2: + if planned is None: return None # nothing substantial to compress + system_msgs, to_summarize, to_keep = planned # Meta-summary: if to_summarize contains multiple existing summary messages # that are long enough to crowd the summarizer input, consolidate them first. @@ -454,9 +485,15 @@ summary_text_input, images = _format_for_summary(to_summarize) - # Truncate oversized input so the summarizer LLM has room to generate output + # Truncate oversized input so the summarizer LLM has room to generate output. + # Keep head + tail, not just the head: to_summarize is chronological, so a + # head-only cut silently dropped the NEWEST summarized messages — the ones + # closest to the kept window, which bridge old context and the recent turn. if len(summary_text_input) > _MAX_SUMMARY_INPUT_CHARS: - summary_text_input = summary_text_input[:_MAX_SUMMARY_INPUT_CHARS] + "\n…[truncated]" + elision = "\n…[middle elided]…\n" + head = _MAX_SUMMARY_INPUT_CHARS * 3 // 4 + tail = _MAX_SUMMARY_INPUT_CHARS - head - len(elision) + summary_text_input = summary_text_input[:head] + elision + summary_text_input[-tail:] system_prompt = _build_summary_system_prompt(profile) prompt = [ @@ -498,7 +535,7 @@ (rough vision-model estimate). """ chars = sum(len(m.content or "") for m in context) - imgs = sum(500 for m in context if m.images) + imgs = sum(500 * len(m.images) for m in context if m.images) return chars // 3 + imgs def __init__(self) -> None: @@ -640,20 +677,17 @@ ) -> bool: """True if ``compress_session`` would actually shrink the context. - Runs only the partition decision (no LLM call), mirroring - ``compress_context``: turn-based, then the aggressive intra-turn - fallback. Also returns True when the context is over the token threshold - even though partition found nothing — the token-budget fallback above - handles that. Used to emit ``CompressionStarted`` only when compression - will really happen (no misleading "compression" status on a no-op). + Runs the same ``_plan_compression`` decision as ``compress_context`` + (no LLM call). Also returns True when the context is over the token + threshold even though partition found nothing — the token-budget + fallback in ``compress_session`` handles that. Used to emit + ``CompressionStarted`` only when compression will really happen + (no misleading "compression" status on a no-op). """ effective_keep_recent = getattr(self._profile, "compression_keep_recent", None) or keep_recent - to_summarize, _ = partition_messages( + if _plan_compression( context, effective_keep_recent, keep_recent_messages=keep_recent_messages - ) - if len(to_summarize) < 2 and keep_recent_messages is not None and keep_recent_messages > 2: - to_summarize, _ = partition_messages(context, effective_keep_recent, keep_recent_messages=2) - if len(to_summarize) >= 2: + ) is not None: return True # Partition no-op: compression can still shrink via the token-budget # fallback — but only if dropping oldest turns actually reduces the @@ -733,13 +767,10 @@ return None # already within budget new_context = system_msgs + to_keep - dropped = len(non_system) - len(to_keep) + # dropped > 0 is guaranteed here: to_keep == non_system returns None above. summary_text = ( "[Context was too large to summarize. Old messages were truncated to prevent " "the model from exceeding its context window. Some earlier details may have been lost.]" - if dropped - else "[Context was over the token limit with too few messages to summarize; " - "it was left intact but the per-message view may be truncated for the LLM.]" ) return new_context, summary_text @@ -836,6 +867,15 @@ self._real_baseline = None await session_store.save(session) + log.info( + "compressor.compressed", + session_id=session_id, + reason=reason, + messages_before=count_before, + messages_after=len(new_context), + context_tokens=session.context_token_count, + ) + # Archive old messages if the hot table exceeds the configured window. if settings.session_messages_window > 0 and session.db_next_sequence > settings.session_messages_window: threshold = session.db_next_sequence - settings.session_messages_window diff --git a/navi/workers/compressor.py b/navi/workers/compressor.py index 085edf5..395dfb1 100644 --- a/navi/workers/compressor.py +++ b/navi/workers/compressor.py @@ -3,9 +3,7 @@ import structlog from navi.config import settings -from navi.core.compressor import compress_context, should_compress, ContextCompressor -from navi.core.events import ContextCompressed -from navi.llm.base import Message +from navi.core.compressor import ContextCompressor, should_compress from .base import Worker, WorkerContext, WorkerResult @@ -16,6 +14,16 @@ """ Compresses session.context when it approaches the token limit. session.messages (full display history) is never modified. + + The token gate uses the real prompt-token count from the last LLM call + (``ctx.context_tokens``); the compression itself is delegated to + ``ContextCompressor.compress_and_save_session`` so the post-turn path gets + the same pipeline as the pre/mid-turn paths: retry on LLM failure, + hard-truncate fallback, the result-None token-budget fallback, the + 65%-target safety net, real-baseline clearing, and message-window + archiving. The worker used to call ``compress_context`` directly with its + own save logic — a second, weaker pipeline that silently no-op'ed whenever + the summarizer LLM failed while the context stayed over the threshold. """ async def run(self, session, ctx: WorkerContext) -> WorkerResult: @@ -27,69 +35,41 @@ settings.context_compression_threshold): return WorkerResult() - count_before = len(session.context) + compressor = ContextCompressor() + compressor.set_profile(ctx.profile) try: - result = await compress_context( - context=session.context, + # Mirror the midturn path: a long autonomous turn (one user + # message + many tool iterations) is a single conversational + # turn, so turn-based compression alone finds nothing to + # summarize. The intra-turn fallback lets the worker compress + # it too — without this the post-turn worker always no-ops for + # the navi_code shape. + event = await compressor.compress_and_save_session( + session=session, + session_store=ctx.session_store, llm=ctx.llm, model=ctx.model, temperature=settings.context_summary_temperature, + session_id=ctx.session_id, + reason="postturn", keep_recent=settings.context_keep_recent, max_tokens=settings.context_summary_max_tokens, - profile=ctx.profile, - # Mirror the midturn path: a long autonomous turn (one user - # message + many tool iterations) is a single conversational - # turn, so turn-based compression alone finds nothing to - # summarize. The intra-turn fallback lets the worker compress - # it too — without this the post-turn worker always no-ops for - # the navi_code shape. keep_recent_messages=max(12, settings.context_keep_recent * 2), ) except Exception: - log.warning("compression_worker.llm_failed", session_id=ctx.session_id, exc_info=True) + # compress_and_save_session already retries and hard-truncates + # internally; reaching here means something unexpected broke + # (store failure etc.) — log and leave the session untouched. + log.warning("compression_worker.failed", session_id=ctx.session_id, exc_info=True) return WorkerResult() - if result is None: + if event is None: return WorkerResult() - new_context, summary_text = result - - # Mark messages that are no longer part of the LLM context - new_context_ids = {id(m) for m in new_context} - for msg in session.messages: - if id(msg) not in new_context_ids and msg.role != "system": - msg.is_context = False - - # The summary returned by the compressor must also live in messages so - # save() writes it to the normalized table, but it is not displayed. - summary_msg = next((m for m in new_context if m.is_summary), None) - if summary_msg and summary_msg not in session.messages: - summary_msg.is_display = False - session.messages.append(summary_msg) - - # UI marker showing that compression happened - session.messages.append(Message( - role="system", - is_compression=True, - is_context=False, - content=summary_text, - )) - - session.context = new_context - session.context_token_count = ContextCompressor.estimate_context_tokens(new_context) - await ctx.session_store.save(session) - log.info( "compression_worker.done", session_id=ctx.session_id, - before=count_before, - after=len(session.context), + before=event.messages_before, + after=event.messages_after, ) - - return WorkerResult(events=[ContextCompressed( - messages_before=count_before, - messages_after=len(session.context), - summary=summary_text, - context_tokens=session.context_token_count, - max_context_tokens=ctx.max_context_tokens, - )]) + return WorkerResult(events=[event]) \ No newline at end of file diff --git a/tests/unit/core/test_compressor.py b/tests/unit/core/test_compressor.py index d238b0f..04e5a6d 100644 --- a/tests/unit/core/test_compressor.py +++ b/tests/unit/core/test_compressor.py @@ -344,6 +344,52 @@ finally: compressor_module._meta_summarize = original_meta + async def test_summary_input_truncated_head_and_tail(self): + """When to_summarize exceeds _MAX_SUMMARY_INPUT_CHARS, the summarizer + input keeps the head (oldest messages) AND the tail (the messages + closest to the kept window). A head-only cut used to silently drop the + newest summarized work — exactly the part that bridges old context to + the verbatim-kept turns.""" + from navi.core.compressor import _MAX_SUMMARY_INPUT_CHARS + + class _RecordingBackend(FakeLLMBackend): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.prompts = [] + + async def complete(self, messages, **kw): + self.prompts.append(messages) + return await super().complete(messages, **kw) + + backend = _RecordingBackend(responses=["Summary"]) + # 20 turns x ~2000 chars -> ~36k chars of to_summarize (turns 18-19 + # kept), comfortably over the 32k input budget. + context = [Message(role="system", content="sys")] + for i in range(20): + filler = "f" * 1990 + if i == 0: + content = "OLDEST_MARKER " + filler + elif i == 17: # the newest turn that gets summarized + content = filler + " NEWEST_SUMMARIZED_MARKER" + else: + content = filler + context.append(Message(role="user", content=content)) + context.append(Message(role="assistant", content=f"ok {i}")) + + new_context, _ = await compress_context( + context=context, + llm=backend, + model="test", + temperature=0.3, + keep_recent=2, + ) + assert new_context is not None + user_prompt = backend.prompts[0][1].content + assert len(user_prompt) <= _MAX_SUMMARY_INPUT_CHARS + 100 + assert "OLDEST_MARKER" in user_prompt # head kept + assert "NEWEST_SUMMARIZED_MARKER" in user_prompt # tail kept + assert "middle elided" in user_prompt # elision marked + async def test_intra_turn_fallback_aggressive(self): """When turn-based partition has nothing to compress but keep_recent_messages is set, an aggressive fallback (keep_recent_messages=2) should still find @@ -391,6 +437,14 @@ context = [Message(role="user", content="hi", images=["base64img"])] assert compressor.estimate_context_tokens(context) == 500 # 2 chars // 3 = 0 + 500 + def test_estimate_context_tokens_counts_each_image(self): + """Images count 500 tokens EACH, not 500 per message that happens to + carry them — multi-image uploads (webclient sends several) used to be + undercounted.""" + compressor = ContextCompressor() + context = [Message(role="user", content="hi", images=["img1", "img2", "img3"])] + assert compressor.estimate_context_tokens(context) == 1500 + @pytest.mark.asyncio async def test_compress_session_success(self): backend = FakeLLMBackend(responses=["Summary text"]) @@ -594,6 +648,33 @@ assert long_result not in text assert "…" in text + def test_format_critical_tool_result_over_budget_keeps_head_and_tail(self): + """A critical result larger than the 4000-char budget keeps head+tail + halves (2000+2000) instead of collapsing to the non-critical 800-char + preview — the all-or-nothing cliff used to throw away 98% of a large + filesystem read while a 3999-char result survived whole.""" + head_marker = "HEAD_MARKER" + "a" * 5000 + tail_marker = "b" * 5000 + "TAIL_MARKER" + long_result = head_marker + "MIDDLE_DROPPED" + tail_marker + msgs = [ + Message( + role="assistant", + tool_calls=[ToolCallRequest(id="1", name="filesystem", arguments={})], + ), + Message( + role="tool", + content=long_result, + name="filesystem", + tool_call_id="1", + is_compression_critical=True, + ), + ] + text, _ = _format_for_summary(msgs) + assert "HEAD_MARKER" in text # head half kept + assert "TAIL_MARKER" in text # tail half kept + assert "MIDDLE_DROPPED" not in text # middle elided + assert "middle elided" in text + def test_summary_prompt_uses_profile_compression_prompt_file(self): profile = make_profile("navi_code", compression_prompt_file="compression_prompt.txt") prompt = _build_summary_system_prompt(profile) diff --git a/tests/unit/workers/test_compression_worker.py b/tests/unit/workers/test_compression_worker.py index a407b32..b100123 100644 --- a/tests/unit/workers/test_compression_worker.py +++ b/tests/unit/workers/test_compression_worker.py @@ -70,6 +70,70 @@ assert result.events == [] +@pytest.mark.asyncio +async def test_worker_hard_truncates_when_summarizer_always_fails(): + """The worker delegates to compress_and_save_session, so a summarizer LLM + that always fails still ends with the hard-truncate fallback — the context + shrinks and ContextCompressed is emitted. The old worker called + compress_context directly and silently no-op'ed on any LLM failure, + leaving the session over the threshold until the next turn's gates.""" + import navi.core.compressor as compressor_module + + worker = CompressionWorker() + backend = FakeLLMBackend(responses=["unused"]) + store = InMemorySessionStore() + session = Session(profile_id="test") + store._sessions[session.id] = session + # 4 turns, each assistant ~20k tokens -> hard-truncate (0.5 * 65536) + # keeps only the newest turn. + big = "x" * 60_000 + for i in range(4): + session.context.append(Message(role="user", content=str(i))) + session.context.append(Message(role="assistant", content=big)) + + async def _always_fail(*args, **kwargs): + raise RuntimeError("summarizer down") + + original = compressor_module.compress_context + compressor_module.compress_context = _always_fail + try: + result = await worker.run(session, _ctx(session.id, backend, store)) + finally: + compressor_module.compress_context = original + + compressed = [ev for ev in result.events if isinstance(ev, ContextCompressed)] + assert compressed, "worker must not silently no-op when the summarizer fails" + assert compressed[0].messages_after < compressed[0].messages_before + assert "truncated" in compressed[0].summary.lower() + assert len(session.context) == compressed[0].messages_after + + +@pytest.mark.asyncio +async def test_worker_marks_dropped_messages_not_in_context(): + """Delegated path keeps the message marking: dropped messages are flagged + is_context=False so a reload does not resurrect them into the context.""" + worker = CompressionWorker() + backend = FakeLLMBackend(responses=["worker summary"]) + store = InMemorySessionStore() + session = Session(profile_id="test") + store._sessions[session.id] = session + session.context.append(Message(role="user", content="1")) + session.context.append(Message(role="assistant", content="a1")) + session.context.append(Message(role="user", content="2")) + session.context.append(Message(role="assistant", content="a2")) + for i in range(12): + session.context.append(Message(role="user", content=f"task {i}")) + session.context.append(Message(role="assistant", content=f"answer {i}")) + session.messages = list(session.context) + + result = await worker.run(session, _ctx(session.id, backend, store)) + assert any(isinstance(ev, ContextCompressed) for ev in result.events) + kept_ids = {id(m) for m in session.context} + dropped = [m for m in session.messages if id(m) not in kept_ids and m.role != "system"] + assert dropped, "old turns must be marked is_context=False" + assert all(m.is_context is False for m in dropped) + + def test_worker_keep_recent_messages_mirrors_midturn(): """The worker passes keep_recent_messages=max(12, context_keep_recent*2), matching the midturn auto-compress path (source of truth for the intra-turn