diff --git a/navi/core/agent.py b/navi/core/agent.py index 22d1d03..a99a86b 100644 --- a/navi/core/agent.py +++ b/navi/core/agent.py @@ -77,6 +77,7 @@ if TYPE_CHECKING: from navi.context_providers._loader import ContextProviderRegistry from navi.memory.store import MemoryStore + from navi.profiles.base import AgentProfile from navi.workers.base import Worker @@ -635,7 +636,7 @@ built_ctx.append(_stall_msg) try: - self._compressor.check_context_size(built_ctx) + self._compressor.check_context_size(built_ctx, session_context=session.context) except ContextTooLargeError as e: # Surface the error as a Navi response (not a raw system error) so the # user sees a coherent message and the exchange is saved to history. @@ -670,6 +671,14 @@ ): yield _ev + # Record the real prompt-token count for the context we just sent so + # the next iteration's midturn/check estimates use real tokens for + # the bulk of the context (not the chars//3 heuristic, which + # undercounts code-heavy tool output and fires compression too late). + # len(session.context) here is exactly what was sent — the new + # assistant/tool messages are appended only below. + self._compressor.record_real_baseline(len(session.context), state.context_tokens) + # Stopped mid-stream — save partial response and exit if stop_event and stop_event.is_set(): if state.accumulated_text: @@ -716,7 +725,7 @@ ) for event in await self._run_workers( - session, llm, profile.model, state.context_tokens + session, llm, profile.model, state.context_tokens, profile=profile ): yield event return @@ -796,6 +805,7 @@ llm: LLMBackend, model: str, context_tokens: int | None, + profile: "AgentProfile | None" = None, ) -> list[AgentEvent]: """Run all workers sequentially; collect their events.""" from navi.workers.base import WorkerContext @@ -808,6 +818,7 @@ model=model, temperature=settings.context_summary_temperature, session_store=self._sessions, + profile=profile, ) events: list[AgentEvent] = [] for worker in self._workers: @@ -891,7 +902,9 @@ extra_system=ctx_injections, session_id=session_id, ) - estimated_tokens = self._compressor.estimate_context_tokens(preflight_ctx) + estimated_tokens = self._compressor.real_baseline_estimate( + session.context, preflight_ctx + ) if should_compress( estimated_tokens, settings.ollama_num_ctx, diff --git a/navi/core/compressor.py b/navi/core/compressor.py index f61843c..f92b1a5 100644 --- a/navi/core/compressor.py +++ b/navi/core/compressor.py @@ -442,11 +442,45 @@ def __init__(self) -> None: self._profile: "AgentProfile | None" = None + # Real prompt-token baseline from the last LLM call: + # (len(session.context) at call time, real prompt_tokens returned). + # The next estimate reuses these real tokens for the bulk of the context + # and only heuristics the messages appended since, instead of the + # chars//3 estimate which undercounts code-heavy tool output and fires + # compression too late. Cleared after compression (context shrank). + self._real_baseline: tuple[int, int] | None = None def set_profile(self, profile: "AgentProfile | None") -> None: """Tell the compressor which profile is active so it can use profile-specific settings.""" self._profile = profile + def record_real_baseline(self, context_len: int, real_tokens: int | None) -> None: + """Cache the real prompt-token count returned by the last LLM call. + + ``context_len`` is ``len(session.context)`` at the time of that call + (before the new assistant/tool messages are appended). The next + estimate reuses these real tokens for the bulk of the context and only + heuristics the new messages appended since. + """ + if real_tokens and real_tokens > 0: + self._real_baseline = (context_len, int(real_tokens)) + + def real_baseline_estimate( + self, session_context: list[Message], built_ctx: list[Message] + ) -> int: + """Estimate context tokens using the real baseline when available. + + Falls back to the chars//3 heuristic on ``built_ctx`` when no baseline + is recorded or it is stale (the session shrank, e.g. right after + compression — ``base_len`` exceeds the current length). + """ + if self._real_baseline is not None: + base_len, base_tokens = self._real_baseline + if base_len <= len(session_context): + delta = self.estimate_context_tokens(session_context[base_len:]) + return base_tokens + delta + return self.estimate_context_tokens(built_ctx) + async def compress_session( self, context: list[Message], @@ -629,7 +663,9 @@ ) return new_context, summary_text - def check_context_size(self, context: list[Message]) -> None: + def check_context_size( + self, context: list[Message], session_context: list[Message] | None = None + ) -> None: """Raise ContextTooLargeError before an LLM call if the context is dangerously large.""" from navi.config import settings from navi.exceptions import ContextTooLargeError @@ -638,7 +674,14 @@ return output_reserve = settings.output_reserve_tokens - total = self.estimate_context_tokens(context) + # When the caller supplies the raw session context, prefer the real-token + # baseline (real bulk + heuristic delta) over the pure chars//3 estimate, + # so a code-heavy context is caught before it overflows the window. + total = ( + self.real_baseline_estimate(session_context, context) + if session_context is not None + else self.estimate_context_tokens(context) + ) available = settings.ollama_num_ctx - output_reserve if total > available: @@ -708,6 +751,9 @@ session.context = new_context session.context_token_count = self.estimate_context_tokens(new_context) + # 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 await session_store.save(session) # Archive old messages if the hot table exceeds the configured window. diff --git a/navi/workers/base.py b/navi/workers/base.py index 0b206a9..c344a7b 100644 --- a/navi/workers/base.py +++ b/navi/workers/base.py @@ -11,6 +11,7 @@ if TYPE_CHECKING: from navi.core.session import Session, SessionStore from navi.llm.base import LLMBackend + from navi.profiles.base import AgentProfile @dataclass @@ -24,6 +25,9 @@ model: str temperature: float session_store: SessionStore + # The active profile so workers can apply profile-specific overrides + # (e.g. compression_keep_recent, compression_max_tokens, prompt file). + profile: "AgentProfile | None" = None @dataclass diff --git a/navi/workers/compressor.py b/navi/workers/compressor.py index 98c518f..085edf5 100644 --- a/navi/workers/compressor.py +++ b/navi/workers/compressor.py @@ -36,6 +36,7 @@ temperature=settings.context_summary_temperature, 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 diff --git a/tests/unit/core/test_compressor.py b/tests/unit/core/test_compressor.py index 6c848cf..904a3c6 100644 --- a/tests/unit/core/test_compressor.py +++ b/tests/unit/core/test_compressor.py @@ -643,3 +643,78 @@ keep_recent=12, keep_recent_messages=16, ) assert result is None + + +class TestRealBaselineEstimate: + """real_baseline_estimate uses real prompt_tokens for the bulk of the + context + a heuristic delta for messages appended since, instead of the + chars//3 estimate that undercounts code-heavy tool output.""" + + def test_no_baseline_falls_back_to_heuristic(self): + c = ContextCompressor() + ctx = [Message(role="user", content="hello world")] # 11 chars -> 3 tokens + assert c.real_baseline_estimate(ctx, ctx) == 3 + + def test_uses_real_baseline_plus_delta_for_new_messages(self): + c = ContextCompressor() + # Simulate: last LLM call saw 4 session messages and reported 5000 real + # prompt tokens. Since then 2 new messages (12 chars -> 4 tokens) were + # appended. Estimate should be 5000 + 4, NOT the chars//3 of the whole. + base = [ + Message(role="user", content="do the task"), + Message(role="assistant", content="step"), + Message(role="tool", name="fs", content="result"), + Message(role="assistant", content="ok"), + ] + c.record_real_baseline(len(base), 5000) + session_context = base + [ + Message(role="tool", name="fs", content="newresult"), + Message(role="assistant", content="done"), + ] + built = [Message(role="system", content="sys")] + session_context + assert c.real_baseline_estimate(session_context, built) == 5004 + + def test_stale_baseline_after_shrink_falls_back_to_heuristic(self): + """After compression the context shrank below baseline_len -> the + baseline is stale and must fall back to heuristic (not crash / neg).""" + c = ContextCompressor() + c.record_real_baseline(10, 5000) # baseline recorded at len 10 + ctx = [Message(role="user", content="hi")] # now len 1 < 10 + # falls back to heuristic on built_ctx + assert c.real_baseline_estimate(ctx, ctx) == 0 # 2 chars // 3 = 0 + + def test_record_real_baseline_ignores_none_and_zero(self): + c = ContextCompressor() + c.record_real_baseline(4, None) + assert c._real_baseline is None + c.record_real_baseline(4, 0) + assert c._real_baseline is None + c.record_real_baseline(4, 5000) + assert c._real_baseline == (4, 5000) + + def test_check_context_size_uses_baseline_when_session_context_given(self): + """check_context_size prefers the real baseline when the caller passes + session_context; without it (subagent / unit tests) it stays heuristic.""" + import navi.config as _config + from navi.config import Settings + + prev_settings = _config.settings + _config.settings = Settings( + _env_file=None, ollama_num_ctx=100_000, output_reserve_tokens=2048, + navi_persona_file="", + ) + try: + c = ContextCompressor() + # Real baseline says we are at 98k tokens; heuristic (chars//3) on a + # tiny string would say ~0. With the baseline the guard must trip + # (98k > 100k - 2048 reserve = 97952 available). + base = [Message(role="user", content="task")] + c.record_real_baseline(len(base), 98_000) + session_context = base + [Message(role="assistant", content="ok")] + built = [Message(role="system", content="sys")] + session_context + from navi.exceptions import ContextTooLargeError + + with pytest.raises(ContextTooLargeError): + c.check_context_size(built, session_context=session_context) + finally: + _config.settings = prev_settings