From 4f852e662fdd53f682188e8d36a8b5022af68277 Mon Sep 17 00:00:00 2001
From: 0xhis <0xhis@users.noreply.github.com>
Date: Sat, 21 Mar 2026 00:44:43 -0700
Subject: [PATCH] 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.
---
 strix/llm/llm.py               | 9 +++++++--
 strix/llm/memory_compressor.py | 9 ++++++++-
 2 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/strix/llm/llm.py b/strix/llm/llm.py
index fe1758f1..969710e1 100644
--- a/strix/llm/llm.py
+++ b/strix/llm/llm.py
@@ -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)
diff --git a/strix/llm/memory_compressor.py b/strix/llm/memory_compressor.py
index 8cad5107..728831b1 100644
--- a/strix/llm/memory_compressor.py
+++ b/strix/llm/memory_compressor.py
@@ -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
         )