From 99b82f48ec7c29343e9f8f29d395f735042436aa Mon Sep 17 00:00:00 2001
From: Soulter <905617992@qq.com>
Date: Mon, 5 Jan 2026 11:34:18 +0800
Subject: [PATCH] feat: enhance context compression with token tracking and
 logging

---
 astrbot/core/agent/context/manager.py | 30 ++++++++++++++++-----------
 1 file changed, 18 insertions(+), 12 deletions(-)

diff --git a/astrbot/core/agent/context/manager.py b/astrbot/core/agent/context/manager.py
index d2ec13f8d..4a133144c 100644
--- a/astrbot/core/agent/context/manager.py
+++ b/astrbot/core/agent/context/manager.py
@@ -65,19 +65,17 @@ class ContextManager:
             # 2. 基于 token 的压缩
             if self.config.max_context_tokens > 0:
                 # check if the messages need to be compressed
-                needs_compression, _ = await self._initial_token_check(result)
+                needs_compression, tokens = await self._initial_token_check(result)
 
                 # compress/truncate the messages if needed
-                result = await self._run_compression(result, needs_compression)
+                result = await self._run_compression(result, tokens, needs_compression)
 
             return result
         except Exception as e:
             logger.error(f"Error during context processing: {e}", exc_info=True)
             return messages
 
-    async def _initial_token_check(
-        self, messages: list[Message]
-    ) -> tuple[bool, int | None]:
+    async def _initial_token_check(self, messages: list[Message]) -> tuple[bool, int]:
         """
         Check if the messages need to be compressed.
 
@@ -88,22 +86,19 @@ class ContextManager:
             tuple: (whether to compress, initial token count)
         """
         if not messages:
-            return False, None
+            return False, 0
         if self.config.max_context_tokens <= 0:
-            return False, None
+            return False, 0
 
         total_tokens = self.token_counter.count_tokens(messages)
 
-        logger.debug(
-            f"ContextManager: total tokens = {total_tokens}, max_context_tokens = {self.config.max_context_tokens}"
-        )
         usage_rate = total_tokens / self.config.max_context_tokens
 
         needs_compression = usage_rate > self.COMPRESSION_THRESHOLD
-        return needs_compression, total_tokens if needs_compression else None
+        return needs_compression, total_tokens if needs_compression else 0
 
     async def _run_compression(
-        self, messages: list[Message], needs_compression: bool
+        self, messages: list[Message], prev_tokens: int, needs_compression: bool
     ) -> list[Message]:
         """
         Compress/truncate the messages if needed.
@@ -120,10 +115,21 @@ class ContextManager:
         if self.config.max_context_tokens <= 0:
             return messages
 
+        logger.debug("Starting context compression process...")
+
         messages = await self.compressor(messages)
 
         # double check
         tokens_after_summary = self.token_counter.count_tokens(messages)
+
+        # calculate compress rate
+        compress_rate = (tokens_after_summary / self.config.max_context_tokens) * 100
+        logger.info(
+            f"Compress completed."
+            f" {prev_tokens} -> {tokens_after_summary} tokens,"
+            f" compression rate: {compress_rate:.2f}%.",
+        )
+
         if (
             tokens_after_summary / self.config.max_context_tokens
             > self.COMPRESSION_THRESHOLD