fix(llm): include system prompt tokens in memory compressor budget

The memory compressor was not accounting for system prompt and agent
identity tokens when calculating the conversation budget. This caused
premature history truncation on long scans with large system prompts.

Adds a reserved_tokens parameter to compress_history() that subtracts
already-accounted tokens from the budget before applying limits.
This commit is contained in:
0xhis 2026-03-21 00:44:43 -07:00
parent c9d2477144
commit 4f852e662f
2 changed files with 15 additions and 3 deletions

View file

@ -10,7 +10,7 @@ from litellm.utils import supports_prompt_caching, supports_vision
from strix.config import Config
from strix.llm.config import LLMConfig
from strix.llm.memory_compressor import MemoryCompressor
from strix.llm.memory_compressor import MemoryCompressor, _get_message_tokens
from strix.llm.utils import (
_truncate_to_first_function,
fix_incomplete_tool_call,
@ -210,7 +210,12 @@ class LLM:
}
)
compressed = list(self.memory_compressor.compress_history(conversation_history))
reserved_tokens = sum(
_get_message_tokens(msg, self.config.litellm_model) for msg in messages
)
compressed = list(
self.memory_compressor.compress_history(conversation_history, reserved_tokens)
)
conversation_history.clear()
conversation_history.extend(compressed)
messages.extend(compressed)

View file

@ -166,9 +166,16 @@ class MemoryCompressor:
def compress_history(
self,
messages: list[dict[str, Any]],
reserved_tokens: int = 0,
) -> list[dict[str, Any]]:
"""Compress conversation history to stay within token limits.
Args:
messages: Conversation history messages to compress.
reserved_tokens: Tokens already reserved for system prompt and
other framing messages outside the conversation history.
Subtracted from the budget before checking limits.
Strategy:
1. Handle image limits first
2. Keep all system messages
@ -201,7 +208,7 @@ class MemoryCompressor:
# Type assertion since we ensure model_name is not None in __init__
model_name: str = self.model_name # type: ignore[assignment]
total_tokens = sum(
total_tokens = reserved_tokens + sum(
_get_message_tokens(msg, model_name) for msg in system_msgs + regular_msgs
)