"""`ContextCompressor` — retry, truncate fallback and persistence around `compress_context`."""
from collections.abc import Callable
from datetime import datetime, timezone
from typing import TYPE_CHECKING
import structlog
from navi.config import settings
from navi.llm.base import LLMBackend, Message
from navi.memory._chunks import build_chunks
from navi.memory.breadcrumbs import MAX_BREADCRUMBS, extract_breadcrumbs
from ..context_builder import truncate_oversized
from ..events import ContextCompressed
from ..image_tokens import context_image_tokens
from ..token_calibration import calibration_for
from .partition import _plan_compression, _split_turn_tail, should_compress
from .summarize import compress_context
if TYPE_CHECKING:
from navi.profiles.base import AgentProfile
log = structlog.get_logger()
# A call that reports fewer real prompt tokens than this share of our own
# estimate of the prompt was sent a truncated prompt (Ollama keeps roughly
# half the window and evaluates that). The estimate errs high on images, so
# the band is deliberately wide — a false positive costs one extra
# compression pass, a missed one costs the session.
_TRUNCATION_RATIO = 0.7
# 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
class ContextCompressor:
"""High-level context compression with retry strategy and hard-truncate fallback.
Thin wrapper around `compress_context` that adds:
1. Retry with keep_recent + 4 on LLM failure.
2. Hard-truncate fallback (drop oldest messages without summarizing).
"""
@staticmethod
def context_chars(context: list[Message]) -> int:
"""Characters of message text — the input to every token estimate."""
return sum(len(m.content or "") for m in context)
@staticmethod
def estimate_context_tokens(context: list[Message]) -> int:
"""Raw chars // 3 heuristic, no calibration.
Kept as the reference unit for tests and for callers that want the
old, deliberately under-counting number. Window decisions must use
``estimate_tokens`` instead — on prod this heuristic landed at ~56k
where the real prompt was 101k, and targets expressed in it were met
while the real prompt sat above the window.
"""
chars = ContextCompressor.context_chars(context)
return chars // 3 + context_image_tokens(context)
@staticmethod
def estimate_tokens(context: list[Message], model=None) -> int:
"""Token estimate in real-prompt units, used by every window decision.
``estimate_context_tokens`` converted to real prompt tokens by the
calibration fitted for ``model`` (see ``navi/core/token_calibration.py``
for why chars // 3 is not usable as-is). Without observations the
calibration falls back to ``DEFAULT_CHARS_PER_TOKEN``.
"""
chars = ContextCompressor.context_chars(context)
return calibration_for(model).tokens_for(chars) + context_image_tokens(context)
@staticmethod
def estimate_tokens_delta(context: list[Message], model=None) -> int:
"""Price messages *appended since* a real baseline, same units.
Uses the calibration slope without its intercept: the fixed prompt
overhead is already part of the baseline count. Pricing a delta with
the raw ``chars // 3`` left the pre/mid-turn gates in one unit and the
window in another for exactly the tool-heavy turns they exist for.
"""
chars = ContextCompressor.context_chars(context)
return calibration_for(model).delta_for(chars) + context_image_tokens(context)
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
# Set when the last call reported far fewer tokens than we sent: the
# backend truncated the prompt. Consumed by the post-turn worker, which
# compresses unconditionally in that case (see record_real_baseline).
self._context_truncated = False
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
@property
def context_truncated(self) -> bool:
"""True when the last recorded call saw its prompt cut by the backend."""
return self._context_truncated
def pop_context_truncated(self) -> bool:
"""Read and clear the truncation flag (the post-turn worker consumes it)."""
value = self._context_truncated
self._context_truncated = False
return value
def record_real_baseline(
self,
context_len: int,
real_tokens: int | None,
*,
chars: int | None = None,
model=None,
sent_tokens: int | None = None,
image_tokens: int = 0,
) -> 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.
``chars`` is the character count of the exact prompt that was sent
(the built context, not just session.context). Together with
``real_tokens`` it feeds the per-model calibration line, which is what
converts hypothetical contexts (the proposed kept region) into real
tokens. Prompts that carried images are kept out of that fit: image
bytes are not in ``chars``, so their tokens would land in the fitted
intercept/slope and skew every later text estimate for that model.
``sent_tokens`` is our estimate of the prompt that was sent. When the
backend reports far fewer tokens than that, the prompt was truncated —
Ollama silently keeps about half the window and reports only what it
evaluated. Caching that number as the baseline used to pin the session
at ~50%: every later estimate started from the truncated count, the
90% trigger never fired, and the model kept receiving a prompt whose
head (system prompt included) had been dropped. A truncated count is
therefore not recorded; the previous — honest — baseline stays, the
estimate keeps its size, and ``context_truncated`` tells the post-turn
worker to compress regardless of the threshold.
"""
if not real_tokens or real_tokens <= 0:
return
if (
sent_tokens
and sent_tokens > 0
and real_tokens < sent_tokens * _TRUNCATION_RATIO
):
self._context_truncated = True
log.warning(
"compressor.prompt_truncated",
sent_tokens=sent_tokens,
real_tokens=real_tokens,
model=model,
)
return
self._real_baseline = (context_len, int(real_tokens))
if chars and not image_tokens:
calibration_for(model).observe(chars, real_tokens)
def real_baseline_estimate(
self,
session_context: list[Message],
built_ctx: list[Message],
model=None,
) -> int:
"""Estimate context tokens using the real baseline when available.
Falls back to the calibrated ``estimate_tokens`` 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). Every
branch is in real prompt tokens, the unit the window is measured in.
Both the delta and the fallback are priced on the **truncated view**
(``truncate_oversized``), i.e. on what ``build()`` actually sends. The
raw stored context may hold a megabyte-long tool result that never
reaches the model; pricing it here made the gates fire for work the
compressor intentionally refuses (a single message cannot be split)
and — worse — made ``check_context_size`` kill the run with "context
too large" over a prompt the model never received.
"""
if self._real_baseline is not None:
base_len, base_tokens = self._real_baseline
if base_len <= len(session_context):
delta = self.estimate_tokens_delta(
truncate_oversized(session_context[base_len:]), model
)
return base_tokens + delta
return self.estimate_tokens(truncate_oversized(built_ctx), model)
async def compress_session(
self,
context: list[Message],
llm: LLMBackend,
model: "list[str] | str | None",
temperature: float,
keep_recent: int,
max_tokens: int | None = None,
keep_recent_messages: int | None = None,
window: int | None = None,
) -> tuple[list[Message], str] | None:
"""Compress context with retry + hard-truncate fallback.
Returns (new_context, summary_text) or None if nothing changed.
Does NOT mutate the session — the caller is responsible for updating
session.context, session.context_token_count, and persisting.
"""
effective_keep_recent = getattr(self._profile, "compression_keep_recent", None) or keep_recent
effective_max_tokens = getattr(self._profile, "compression_max_tokens", None) or max_tokens
window_tokens = window or settings.ollama_num_ctx
# Every size decision in this pass — the keep-region target, the
# post-check, the truncate budgets — must be in real-token units, the
# same units the trigger fired in, AND on the view the model is sent
# (``truncate_oversized``): pricing raw stored messages made a kept
# megabyte-long tool result look like the whole window and drop turns
# the model could have had. Mixed units are what let a compression
# pass finish at "65% of the window" while the real prompt it produced
# was above the window.
est = lambda msgs: self.estimate_tokens(truncate_oversized(msgs), model)
# Attempt 1: normal compression
try:
result = await compress_context(
context=context,
llm=llm,
model=model,
temperature=temperature,
keep_recent=effective_keep_recent,
max_tokens=effective_max_tokens,
keep_recent_messages=keep_recent_messages,
profile=self._profile,
window=window_tokens,
estimate=est,
)
except Exception:
# Attempt 2: keep more recent turns verbatim
try:
result = await compress_context(
context=context,
llm=llm,
model=model,
temperature=temperature,
keep_recent=effective_keep_recent + 4,
max_tokens=effective_max_tokens,
keep_recent_messages=(keep_recent_messages + 4)
if keep_recent_messages is not None
else None,
profile=self._profile,
window=window_tokens,
estimate=est,
)
except Exception:
# Attempt 3: hard-truncate fallback
return self._hard_truncate(context, window=window_tokens, estimate=est)
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.)
if should_compress(
est(context),
window_tokens,
settings.context_compression_threshold,
):
return self._token_budget_truncate(
context, estimate=est, window=window_tokens
)
return None
# Target hysteresis safety net: if the compressed context is still over
# the target fraction (midturn mode kept a huge in-flight turn, or even
# the keep_recent floor could not fit the window), drop oldest kept
# turns by token budget so a compression firing always leaves real
# headroom below the trigger — otherwise the next few messages re-fire
# compression and the context yo-yos at the trigger line. The dropped
# turns were kept verbatim (not summarized) and are lost; this only
# triggers when the verbatim kept region alone exceeds the target.
new_context, summary_text = result
target_tokens = int(window_tokens * settings.context_compression_target)
if est(new_context) > target_tokens:
truncated = self._token_budget_truncate(
new_context, budget_tokens=target_tokens, estimate=est, window=window_tokens
)
if truncated is not None:
new_context = truncated[0]
return new_context, summary_text
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,
current_tokens: int | None = None,
model=None,
) -> bool:
"""True if ``compress_session`` would actually shrink the context.
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
window_tokens = max_context_tokens or settings.ollama_num_ctx
# Same unit as the caller's real-token estimate: the truncated view the
# model is actually sent, not the raw stored messages.
est = lambda msgs: self.estimate_tokens(truncate_oversized(msgs), model)
if _plan_compression(
context,
effective_keep_recent,
keep_recent_messages=keep_recent_messages,
window=window_tokens,
estimate=est,
) 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
# 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.
# The size compared against the threshold is the caller's real-token
# estimate (`current_tokens`), not a fresh heuristic: the trigger fired
# on real tokens, and judging it again in chars // 3 units made the two
# disagree by ~1.8x.
if threshold is not None and max_context_tokens is not None:
total = current_tokens if current_tokens is not None else est(context)
if should_compress(total, window_tokens, threshold):
return self._token_budget_truncate(
context, estimate=est, window=window_tokens
) is not None
return False
def _hard_truncate(
self,
context: list[Message],
window: int | None = None,
estimate: Callable[[list[Message]], int] | None = None,
) -> 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.
"""
budget = max(1, int((window or settings.ollama_num_ctx) * _HARD_TRUNCATE_TOKEN_FRAC))
return self._token_budget_truncate(
context, budget_tokens=budget, estimate=estimate, window=window
)
def _token_budget_truncate(
self,
context: list[Message],
budget_tokens: int | None = None,
estimate: Callable[[list[Message]], int] | None = None,
window: 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.
"""
est = estimate or ContextCompressor.estimate_context_tokens
if budget_tokens is None:
budget_tokens = max(
1, int((window or 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"]
if not non_system:
return None
# Group non-system into turns (user message starts a turn)
turns: list[list[Message]] = []
current: list[Message] = []
for msg in non_system:
if msg.role == "user" and current:
turns.append(current)
current = [msg]
else:
current.append(msg)
if current:
turns.append(current)
# 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]] = []
used = 0
for turn in reversed(turns):
turn_tokens = est(turn)
if kept_turns and used + turn_tokens > budget_tokens:
break
kept_turns.insert(0, turn)
used += turn_tokens
if kept_turns and used > budget_tokens:
# The newest turn alone is over budget (a first turn of a few huge
# messages): the loop above always keeps it whole, which used to
# make this fallback a no-op exactly when it was needed. Trim inside
# it — the user head and the newest messages stay verbatim.
newest = kept_turns[-1]
split = _split_turn_tail(
newest,
len(newest),
budget_tokens=budget_tokens,
estimate=est,
)
if split is not None:
kept_turns[-1] = split[1]
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 > 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.]"
)
return new_context, summary_text
def check_context_size(
self,
context: list[Message],
session_context: list[Message] | None = None,
model=None,
window: int | None = None,
) -> None:
"""Raise ContextTooLargeError before an LLM call if the context is dangerously large."""
# Local import: the window and the reserve are read at call time, so a
# runtime settings swap (tests, operator reload) reaches this check.
from navi.config import settings as _settings
from navi.exceptions import ContextTooLargeError
if not context:
return
window_tokens = window or _settings.ollama_num_ctx
# A reserve bigger than the window turns `available` negative and makes
# every check fail, however small the context — a 500-token model still
# has to be able to answer. Cap the reserve at a quarter of the window.
output_reserve = min(
_settings.output_reserve_tokens, max(1, window_tokens // 4)
)
est = lambda msgs: self.estimate_tokens(msgs, model)
# When the caller supplies the raw session context, prefer the real-token
# baseline (real bulk + calibrated delta) over the calibrated estimate;
# both are in the same units as the real counts the window is measured in.
total = (
self.real_baseline_estimate(session_context, context, model)
if session_context is not None
else est(context)
)
available = window_tokens - output_reserve
if total > available:
existing = est(context[:-1])
new = est(context[-1:])
remaining = available - existing
raise ContextTooLargeError(
f"Context too large: new content is ~{new:,} estimated tokens, "
f"but only ~{max(0, remaining):,} tokens are available "
f"(window {window_tokens:,}, already used ~{existing:,}, "
f"output_reserve {output_reserve:,}). "
"Split the file into smaller parts or delegate to a subagent."
)
async def compress_and_save_session(
self,
session,
session_store,
llm: LLMBackend,
model: str,
temperature: float,
session_id: str,
reason: str,
keep_recent: int,
max_tokens: int | None = None,
keep_recent_messages: int | None = None,
window: int | None = None,
memory_store=None,
) -> ContextCompressed | None:
"""Compresses the session context and persists the changes to the session store."""
count_before = len(session.context)
window_tokens = window or settings.ollama_num_ctx
result = await self.compress_session(
context=session.context,
llm=llm,
model=model,
temperature=temperature,
keep_recent=keep_recent,
max_tokens=max_tokens,
keep_recent_messages=keep_recent_messages,
window=window_tokens,
)
if result is None:
return None
new_context, summary_text = result
# Mark messages in session.messages as not context if they are no longer in new_context
# and are not system messages.
# Iterate session.context too: after a session reload, context-only rows
# (old summaries — is_display=False) are present in session.context but
# NOT in session.messages, so demoting only the latter left stale
# summaries with is_context=True in the DB — they piled up in the
# context at every reload (the half-the-context-is-summaries symptom).
new_context_ids = {id(m) for m in new_context}
demoted: list[Message] = []
seen: set[int] = set()
for msg in [*session.messages, *session.context]:
if msg.role == "system" or id(msg) in new_context_ids or id(msg) in seen:
continue
seen.add(id(msg))
if not msg.is_context:
# Dropped by an earlier compression. Iterating the display
# history surfaces it again on every later pass; indexing it
# again writes the same turns under a new chunk boundary (new
# hash) instead of once, forever. The current summary covers
# it; if it was dropped while indexing was off, that is the
# one turn range this trade-off leaves unindexed.
continue
msg.is_context = False
demoted.append(msg)
# Rows that live only in session.context — a summary is is_display=False,
# so after a reload it is absent from session.messages — are invisible to
# save(), which updates flags for display history only. Their demotion
# used to stay in memory: every reload brought the stale summary back
# with is_context=True, and one prod session ended up carrying 15 live
# summaries, 80% of its context. Persist those flips explicitly.
in_messages = {id(m) for m in session.messages}
context_only = [
m.sequence_number
for m in demoted
if m.sequence_number >= 0 and id(m) not in in_messages
]
if context_only:
await session_store.set_context_flag(session_id, context_only, False)
# Add the summary message to session.messages if it's not already there
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)
# Keep the dropped turns searchable before they leave the context for
# good. This must happen while `demoted` is still in hand: after the
# save there is no object that says "these were the turns compression
# removed", only the display history.
await self._index_demoted(memory_store, session, session_id, demoted, summary_msg)
# Add a system message with is_compression=True and content=summary_text
session.messages.append(
Message(
role="system",
content=summary_text,
is_compression=True,
is_context=False,
created_at=datetime.now(timezone.utc),
)
)
session.context = new_context
# 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. The truncation
# flag goes with it: the prompt that will be sent is small again.
self._real_baseline = None
self._context_truncated = False
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
archived = await session_store.archive_old_messages(session_id, threshold)
if archived > 0:
session.messages = [m for m in session.messages if m.sequence_number >= threshold]
session.context = [m for m in session.context if m.sequence_number >= threshold]
session.archive_threshold = threshold
return ContextCompressed(
messages_before=count_before,
messages_after=len(new_context),
summary=summary_text,
context_tokens=session.context_token_count,
max_context_tokens=window_tokens,
)
async def _index_demoted(
self,
memory_store,
session,
session_id: str,
demoted: list[Message],
summary_msg: Message | None,
) -> None:
"""Write the turns compression just dropped into the raw-chunk index.
Append-only: re-running compression on the same history hits the
`(session_id, content_hash)` uniqueness and writes nothing new.
Breadcrumbs come from the summary's `KEYS:` line and, always, from the
deterministic extractor over the chunk text. The hard-truncate fallback
never calls a summarizer, so its index would otherwise be empty; and the
extractor also catches names the summarizer chose not to list. Every
chunk of this pass carries the union — the keys describe the batch, and
a searcher does not know which chunk of the batch holds the name.
"""
if memory_store is None or not demoted or not settings.memory_search_enabled:
return
try:
chunks = build_chunks(demoted)
if not chunks:
return
summary_keys = []
if summary_msg is not None:
summary_keys = list(
(summary_msg.metadata or {}).get("breadcrumbs") or []
)
keys = _merge_breadcrumbs(
summary_keys,
extract_breadcrumbs(
[c["content"] for c in chunks],
limit=max(0, MAX_BREADCRUMBS - len(summary_keys)),
),
)
for chunk in chunks:
chunk["breadcrumbs"] = keys
await memory_store.add_chunks(
chunks, user_id=getattr(session, "user_id", None), session_id=session_id
)
except Exception:
# Indexing must never break compression — a missing chunk is a lost
# search result, a raised exception here is a lost turn.
log.warning("compressor.chunk_index_failed", session_id=session_id, exc_info=True)
def _merge_breadcrumbs(*groups: list[str]) -> list[str]:
merged: list[str] = []
seen: set[str] = set()
for group in groups:
for item in group:
key = " ".join(str(item).lower().split())
if not key or key in seen:
continue
seen.add(key)
merged.append(str(item))
if len(merged) >= MAX_BREADCRUMBS:
return merged
return merged
