diff --git a/navi/config.py b/navi/config.py index 53c2411..095b08d 100644 --- a/navi/config.py +++ b/navi/config.py @@ -124,6 +124,11 @@ context_summary_temperature: float = 0.3 context_summary_max_tokens: int = 4000 # max output tokens for the summary LLM call output_reserve_tokens: int = 2048 # headroom reserved for model response in context checks + # Per-message token budget for the LLM context view. A single tool/assistant + # message whose estimated size exceeds this is head/tail-truncated in the + # built context (never in stored history) so one huge tool result cannot + # alone blow the window. 0 = auto (ollama_num_ctx // 6). + context_message_token_budget: int = 0 # Global personality prompt prepended to every agent's system prompt. # Multi-line values don't survive .env parsing reliably, so prefer diff --git a/navi/core/agent.py b/navi/core/agent.py index d63ce21..22d1d03 100644 --- a/navi/core/agent.py +++ b/navi/core/agent.py @@ -849,6 +849,19 @@ settings.context_compression_threshold, ) ): + # Confirm there is actually something to compress before announcing + # it. The token gate can fire for a context that the partition cannot + # shrink (e.g. a few very large messages in a single turn — handled + # instead by per-message view truncation in context_builder.build). + # Emitting CompressionStarted without a following ContextCompressed + # would show a misleading "compression" status with no work done. + if not self._compressor.would_compress( + session.context, + keep_recent=settings.context_keep_recent, + threshold=settings.context_compression_threshold, + max_context_tokens=settings.ollama_num_ctx, + ): + return yield CompressionStarted( context_tokens=self._compressor.estimate_context_tokens(session.context), max_context_tokens=settings.ollama_num_ctx, @@ -884,6 +897,20 @@ settings.ollama_num_ctx, settings.context_compression_threshold, ): + # See _compression_events_preturn: only announce compression when + # the partition (or token-budget fallback) can actually shrink the + # stored context. The per-message view truncation in build() + # already caps single huge messages for the LLM, so a no-op here + # no longer risks an overflow — but announcing it would still lie + # to the UI. + if not self._compressor.would_compress( + session.context, + keep_recent=settings.context_keep_recent, + keep_recent_messages=max(12, settings.context_keep_recent * 2), + threshold=settings.context_compression_threshold, + max_context_tokens=settings.ollama_num_ctx, + ): + return yield CompressionStarted( context_tokens=estimated_tokens, max_context_tokens=settings.ollama_num_ctx, diff --git a/navi/core/compressor.py b/navi/core/compressor.py index 03a9a12..f61843c 100644 --- a/navi/core/compressor.py +++ b/navi/core/compressor.py @@ -258,6 +258,10 @@ # Prevents the summarizer from receiving near-context-sized input it can't fit alongside output. _MAX_SUMMARY_INPUT_CHARS = 24_000 +# Hard-truncate fallbacks keep at most this fraction of the context window's +# tokens of recent messages (the rest is dropped without summarizing). +_HARD_TRUNCATE_TOKEN_FRAC = 0.5 + # 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 @@ -494,21 +498,96 @@ return self._hard_truncate(context) if result is None: + # Partition found nothing to summarize (e.g. few very large messages + # in a single turn). If the context is nonetheless over the token + # threshold, fall back to a token-budget hard-truncate so the gate + # firing always results in shrinkage — otherwise the agent runs until + # the window overflows. (The per-message truncation in + # context_builder.build caps single huge messages for the LLM view; + # this is the stored-context safety net.) + from navi.config import settings + + if should_compress( + self.estimate_context_tokens(context), + settings.ollama_num_ctx, + settings.context_compression_threshold, + ): + return self._token_budget_truncate(context) return None return result + def would_compress( + self, + context: list[Message], + keep_recent: int, + keep_recent_messages: int | None = None, + threshold: float | None = None, + max_context_tokens: int | None = None, + ) -> 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). + """ + effective_keep_recent = getattr(self._profile, "compression_keep_recent", None) or keep_recent + to_summarize, _ = partition_messages( + 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: + return True + # Partition no-op: compression can still shrink via the token-budget + # fallback — but only if dropping oldest turns actually reduces the + # context (a single huge turn that alone exceeds the budget cannot be + # split by message-boundary truncation, and is handled instead by the + # per-message view truncation in context_builder.build). Mirror that + # decision so would_compress and compress_session agree exactly. + if threshold is not None and max_context_tokens is not None: + if should_compress( + self.estimate_context_tokens(context), max_context_tokens, threshold + ): + return self._token_budget_truncate(context) is not None + return False + def _hard_truncate( self, context: list[Message] ) -> tuple[list[Message], str] | None: """Last-resort fallback: drop oldest non-system messages without summarizing. Keeps whole conversational turns so tool call groups are never split. + Bounded by a token budget (half the context window) rather than a fixed + message count, so it still shrinks a context made of a few very large + messages — the fixed-count floor used to no-op on <= 6 messages even + when they were huge. """ + from navi.config import settings + + budget = max(1, int(settings.ollama_num_ctx * _HARD_TRUNCATE_TOKEN_FRAC)) + return self._token_budget_truncate(context, budget_tokens=budget) + + def _token_budget_truncate( + self, context: list[Message], budget_tokens: int | None = None + ) -> tuple[list[Message], str] | None: + """Drop oldest non-system messages (by turn groups) until under budget. + + Keeps system messages and the newest conversational turns whole (tool + call groups stay together). Returns None if the context already fits the + budget or there is nothing to drop. ``budget_tokens`` defaults to half + the context window. + """ + from navi.config import settings + + if budget_tokens is None: + budget_tokens = max(1, int(settings.ollama_num_ctx * _HARD_TRUNCATE_TOKEN_FRAC)) + system_msgs = [m for m in context if m.role == "system"] non_system = [m for m in context if m.role != "system"] - - _HARD_TRUNCATE_KEEP = 6 - if len(non_system) <= _HARD_TRUNCATE_KEEP: + if not non_system: return None # Group non-system into turns (user message starts a turn) @@ -523,20 +602,30 @@ if current: turns.append(current) - # Keep at least 2 recent turns, or enough messages to exceed _HARD_TRUNCATE_KEEP + # Keep newest turns (whole) until the token budget is exhausted; always + # keep at least the most recent turn so the agent retains the in-flight + # request + its latest tool results. kept_turns: list[list[Message]] = [] - kept_count = 0 + used = 0 for turn in reversed(turns): - kept_turns.insert(0, turn) - kept_count += len(turn) - if kept_count >= _HARD_TRUNCATE_KEEP and len(kept_turns) >= 2: + turn_tokens = self.estimate_context_tokens(turn) + if kept_turns and used + turn_tokens > budget_tokens: break + kept_turns.insert(0, turn) + used += turn_tokens to_keep = [m for turn in kept_turns for m in turn] + if len(to_keep) == len(non_system): + return None # already within budget + new_context = system_msgs + to_keep + dropped = len(non_system) - len(to_keep) 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 diff --git a/navi/core/context_builder.py b/navi/core/context_builder.py index 0f549fe..60271cc 100644 --- a/navi/core/context_builder.py +++ b/navi/core/context_builder.py @@ -15,6 +15,28 @@ log = structlog.get_logger() +# Roles whose content may be head/tail-truncated when a single message is too +# large for the context window. User messages carry intent and are never +# truncated; system/summary messages are already compact. +_TRUNCATABLE_ROLES = {"tool", "assistant"} + + +def _truncate_content(content: str, budget_tokens: int) -> str: + """Head/tail-truncate ``content`` to roughly ``budget_tokens`` tokens. + + Keeps the head and tail verbatim with a marker naming the elided middle, so + the model retains the beginning and end of a large tool result (usually the + most informative parts) without the whole thing blowing the context window. + """ + budget_chars = budget_tokens * 3 + if len(content) <= budget_chars: + return content + half = budget_chars // 2 + dropped_tokens = (len(content) - 2 * half) // 3 + marker = f"\n[…truncated ~{dropped_tokens:,} tokens…]\n" + return content[:half] + marker + content[-half:] + + 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`. @@ -335,6 +357,29 @@ lines.append(text) 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 + def build( self, session_context: list[Message], @@ -370,6 +415,7 @@ ) system_msg = Message(role="system", content=system_prompt) conv = [m for m in session_context if m.role != "system"] + conv = self._truncate_oversized(conv) result: list[Message] = [system_msg] if mem: result.append(mem) diff --git a/navi/workers/compressor.py b/navi/workers/compressor.py index 2ac49ec..98c518f 100644 --- a/navi/workers/compressor.py +++ b/navi/workers/compressor.py @@ -36,6 +36,13 @@ temperature=settings.context_summary_temperature, keep_recent=settings.context_keep_recent, max_tokens=settings.context_summary_max_tokens, + # 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) diff --git a/tests/unit/core/test_agent.py b/tests/unit/core/test_agent.py index b454b4a..e10661e 100644 --- a/tests/unit/core/test_agent.py +++ b/tests/unit/core/test_agent.py @@ -411,6 +411,38 @@ async for _ in agent.compact_stream("no-such-session"): pass + @pytest.mark.asyncio + async def test_midturn_no_compression_started_when_nothing_to_compress(self, agent, session, monkeypatch): + """C: even when the token gate fires, CompressionStarted is NOT emitted + if would_compress says the partition cannot shrink the context — no + misleading "compression" status with no work done. The gate is forced + True here so we exercise the would_compress guard directly.""" + import navi.core.agent as agent_mod + from navi.core.events import CompressionStarted + + # Small context that partition cannot shrink (would_compress False). + session.context.append(Message(role="user", content="hi")) + session.context.append(Message(role="assistant", content="hello")) + await agent._sessions.save(session) + + # Force the token gate open; without this the small context would not + # reach the would_compress check at all. + monkeypatch.setattr(agent_mod, "should_compress", lambda *a, **k: True) + + events = [ + ev + async for ev in agent._compression_events_midturn( + session, + llm=agent._get_backend("ollama"), + profile=agent._profiles.get("test"), + session_id=session.id, + iteration=1, + ctx_injections=[], + mem=None, + ) + ] + assert not any(isinstance(ev, CompressionStarted) for ev in events) + # ─── run_ephemeral() tests ─────────────────────────────────────────────────── diff --git a/tests/unit/core/test_compressor.py b/tests/unit/core/test_compressor.py index b17bc11..6c848cf 100644 --- a/tests/unit/core/test_compressor.py +++ b/tests/unit/core/test_compressor.py @@ -385,7 +385,6 @@ async def test_compress_session_retry_on_failure(self): """First attempt fails, second succeeds with keep_recent + 4.""" import navi.core.compressor as compressor_module - from unittest.mock import AsyncMock backend = FakeLLMBackend(responses=["Summary text"]) original_compress_context = compressor_module.compress_context @@ -419,7 +418,7 @@ @pytest.mark.asyncio async def test_compress_session_hard_truncate_on_double_failure(self): - """Both attempts fail → hard-truncate fallback.""" + """Both attempts fail → token-budget hard-truncate drops oldest turns.""" import navi.core.compressor as compressor_module from unittest.mock import AsyncMock @@ -431,16 +430,16 @@ compressor_module.compress_context = _always_fail try: compressor = ContextCompressor() + # Four turns, each large enough that the whole context far exceeds + # the hard-truncate token budget (num_ctx * 0.5). The fallback must + # keep only the newest turn(s) that fit the budget and drop the rest. + big = "x" * 120_000 # ~40k tokens each — two already overflow 0.5*65k context = [ Message(role="system", content="sys"), - Message(role="user", content="1"), - Message(role="assistant", content="a1"), - Message(role="user", content="2"), - Message(role="assistant", content="a2"), - Message(role="user", content="3"), - Message(role="assistant", content="a3"), - Message(role="user", content="4"), - Message(role="assistant", content="a4"), + Message(role="user", content="1"), Message(role="assistant", content=big), + Message(role="user", content="2"), Message(role="assistant", content=big), + Message(role="user", content="3"), Message(role="assistant", content=big), + Message(role="user", content="4"), Message(role="assistant", content=big), ] result = await compressor.compress_session( context=context, @@ -452,8 +451,11 @@ assert result is not None new_context, summary = result assert "truncated" in summary.lower() - # system + 6 kept = 7 - assert len(new_context) == 7 + # System + only the newest turn (user "4" + its assistant) fit the + # budget; the three older turns are dropped. + assert len(new_context) == 3 + assert new_context[0].role == "system" + assert new_context[1].content == "4" finally: compressor_module.compress_context = original_compress_context @@ -548,3 +550,96 @@ # max_tokens passed to llm.complete should be overridden by profile # FakeLLMBackend ignores max_tokens, but we can at least verify the call ran assert backend._call_idx == 1 + + +class TestWouldCompress: + """would_compress must predict compress_session's actual outcome (no LLM call), + so the agent only announces compression when it will really happen.""" + + def test_false_for_small_context(self): + c = ContextCompressor() + ctx = [Message(role="user", content="hi"), Message(role="assistant", content="hello")] + assert c.would_compress(ctx, keep_recent=8, threshold=0.70, max_context_tokens=65536) is False + + def test_true_when_partition_finds_messages(self): + c = ContextCompressor() + ctx = [Message(role="user", content="do task")] + [ + Message(role="assistant", content=f"s{i}") for i in range(20) + ] + # many messages in one turn -> intra-turn split finds to_summarize + assert c.would_compress( + ctx, keep_recent=8, keep_recent_messages=16, threshold=0.70, max_context_tokens=65536 + ) is True + + def test_false_for_single_huge_turn_that_cannot_be_split(self): + """A few very large messages in one turn: the token gate would fire, but + neither partition nor the token-budget truncate can shrink a single turn + that alone exceeds the budget — so would_compress is False (the + per-message view truncation in context_builder.build handles this case).""" + c = ContextCompressor() + big = "x" * 200_000 + ctx = [ + Message(role="user", content="read file"), + Message(role="assistant", content="a"), + Message(role="tool", name="fs", content=big), + ] + assert c.would_compress( + ctx, keep_recent=12, keep_recent_messages=16, threshold=0.70, max_context_tokens=65536 + ) is False + + def test_true_when_token_budget_truncate_can_drop_oldest_turns(self): + """Multiple huge turns: dropping oldest turns brings the context under + budget, so the token-budget fallback will shrink it -> would_compress True.""" + c = ContextCompressor() + big = "x" * 120_000 + ctx = [] + for i in range(4): + ctx.append(Message(role="user", content=str(i))) + ctx.append(Message(role="assistant", content=big)) + assert c.would_compress( + ctx, keep_recent=12, keep_recent_messages=16, threshold=0.70, max_context_tokens=65536 + ) is True + + +class TestTokenBudgetFallback: + """compress_session shrinks via token-budget truncate when partition no-ops + but the context is over the threshold (many large turns).""" + + @pytest.mark.asyncio + async def test_compress_session_drops_oldest_turns_via_token_budget(self): + c = ContextCompressor() + big = "x" * 120_000 # ~40k tokens each + ctx = [Message(role="system", content="sys")] + for i in range(4): + ctx.append(Message(role="user", content=str(i))) + ctx.append(Message(role="assistant", content=big)) + # llm=None is fine: partition no-ops (1 turn? no — 4 turns, but keep_recent=12 + # > 4 turns -> turn-based no-op; intra-turn krm=16 on the LAST turn only has + # 2 msgs -> no-op too). So the token-budget fallback path is taken. + result = await c.compress_session( + context=ctx, llm=None, model="m", temperature=0.3, + keep_recent=12, keep_recent_messages=16, + ) + assert result is not None + new_context, summary = result + # system + only the newest turn fit the budget + assert new_context[0].role == "system" + assert len(new_context) == 3 + assert new_context[1].content == "3" + + @pytest.mark.asyncio + async def test_compress_session_returns_none_when_view_truncation_suffices(self): + """Single huge turn: token-budget truncate can't split it -> None. + (context_builder.build truncates the per-message view instead.)""" + c = ContextCompressor() + big = "x" * 200_000 + ctx = [ + Message(role="user", content="read file"), + Message(role="assistant", content="a"), + Message(role="tool", name="fs", content=big), + ] + result = await c.compress_session( + context=ctx, llm=None, model="m", temperature=0.3, + keep_recent=12, keep_recent_messages=16, + ) + assert result is None diff --git a/tests/unit/core/test_context_builder.py b/tests/unit/core/test_context_builder.py index f583610..462a5f6 100644 --- a/tests/unit/core/test_context_builder.py +++ b/tests/unit/core/test_context_builder.py @@ -255,3 +255,67 @@ ) assert msg is not None assert "navi-1" in msg.content # without cwd we cannot judge scope + + +class TestTruncateOversized: + """Per-message head/tail truncation in build() — caps a single huge tool + result so it cannot alone blow the context window, without mutating stored + history.""" + + def test_truncates_oversized_tool_message(self): + builder = ContextBuilder(profile_registry=make_profile_registry()) + profile = make_profile("test") + big = "LINE\n" * 40_000 # ~200k chars -> well over the per-message budget + ctx = [ + Message(role="user", content="read file"), + Message(role="assistant", content="reading"), + Message(role="tool", name="fs", content=big), + ] + result = builder.build(ctx, profile, mem=None) + tool_msgs = [m for m in result if m.role == "tool"] + assert len(tool_msgs) == 1 + truncated = tool_msgs[0] + assert "truncated" in truncated.content + # head and tail of the original content survive + assert truncated.content.startswith("LINE\n") + assert truncated.content.endswith("LINE\n") + # shrunk to roughly the budget (num_ctx // 6 tokens -> *3 chars) + assert len(truncated.content) < len(big) // 2 + + def test_does_not_mutate_stored_message(self): + builder = ContextBuilder(profile_registry=make_profile_registry()) + profile = make_profile("test") + big = "LINE\n" * 40_000 + original = Message(role="tool", name="fs", content=big) + ctx = [Message(role="user", content="read file"), original] + builder.build(ctx, profile, mem=None) + # The stored message object is untouched — only the LLM view was truncated. + assert len(original.content) == len(big) + assert original.content == big + + def test_user_messages_never_truncated(self): + builder = ContextBuilder(profile_registry=make_profile_registry()) + profile = make_profile("test") + huge = "x" * 200_000 + ctx = [ + Message(role="user", content=huge), + Message(role="assistant", content=huge), + ] + result = builder.build(ctx, profile, mem=None) + user_msgs = [m for m in result if m.role == "user"] + asst = [m for m in result if m.role == "assistant"] + assert len(user_msgs[0].content) == len(huge) # user intent preserved verbatim + assert "truncated" in asst[0].content # assistant IS truncatable + + def test_small_messages_left_intact(self): + builder = ContextBuilder(profile_registry=make_profile_registry()) + profile = make_profile("test") + ctx = [ + Message(role="user", content="hi"), + Message(role="assistant", content="small reply"), + Message(role="tool", name="fs", content="tiny result"), + ] + result = builder.build(ctx, profile, mem=None) + tool_msgs = [m for m in result if m.role == "tool"] + assert tool_msgs[0].content == "tiny result" + assert "truncated" not in tool_msgs[0].content diff --git a/tests/unit/workers/test_compression_worker.py b/tests/unit/workers/test_compression_worker.py new file mode 100644 index 0000000..11f69ae --- /dev/null +++ b/tests/unit/workers/test_compression_worker.py @@ -0,0 +1,81 @@ +"""Tests for the post-turn CompressionWorker.""" + +from __future__ import annotations + +import pytest + +from navi.core.events import ContextCompressed +from navi.core.session import InMemorySessionStore, Session +from navi.workers.base import WorkerContext +from navi.workers.compressor import CompressionWorker +from navi.llm.base import Message +from tests.conftest_factory import FakeLLMBackend + + +def _ctx(session_id: str, llm, store) -> WorkerContext: + return WorkerContext( + session_id=session_id, + context_tokens=50_000, # over the 0.70 threshold -> worker attempts compression + max_context_tokens=65536, + llm=llm, + model="test", + temperature=0.3, + session_store=store, + ) + + +@pytest.mark.asyncio +async def test_worker_compresses_single_long_turn(): + """Regression: a single long autonomous turn (one user message + many tool + iterations = one turn) used to no-op because the worker called + compress_context without keep_recent_messages. Now it mirrors the midturn + path and compresses via the intra-turn fallback.""" + worker = CompressionWorker() + backend = FakeLLMBackend(responses=["summary of the work so far"]) + store = InMemorySessionStore() + session = Session(profile_id="test") + store._sessions[session.id] = session + session.context.append(Message(role="user", content="do the task")) + for i in range(20): + session.context.append(Message(role="assistant", content=f"step {i} " * 40)) + + result = await worker.run(session, _ctx(session.id, backend, store)) + assert any(isinstance(ev, ContextCompressed) for ev in result.events) + compressed = next(ev for ev in result.events if isinstance(ev, ContextCompressed)) + assert compressed.messages_after < compressed.messages_before + assert compressed.summary == "summary of the work so far" + + +@pytest.mark.asyncio +async def test_worker_no_op_when_context_small(): + """Below the token threshold the worker does nothing.""" + worker = CompressionWorker() + backend = FakeLLMBackend(responses=["irrelevant"]) + store = InMemorySessionStore() + session = Session(profile_id="test") + store._sessions[session.id] = session + session.context.append(Message(role="user", content="hi")) + session.context.append(Message(role="assistant", content="hello")) + + ctx = WorkerContext( + session_id=session.id, + context_tokens=10, # well below threshold + max_context_tokens=65536, + llm=backend, + model="test", + temperature=0.3, + session_store=store, + ) + result = await worker.run(session, ctx) + assert result.events == [] + + +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 + fallback). Guards against regressing back to keep_recent_messages=None.""" + import inspect + + src = inspect.getsource(CompressionWorker.run) + assert "keep_recent_messages" in src + assert "max(12" in src \ No newline at end of file