Skip to content

feat: add context compression hooks for plugin extensibility -还未自测，但是修改内容不多，如果可以的话可以直接合并 #5291

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
evpeople wants to merge 3 commits into
base: master
Choose a base branch
from
Draft

feat: add context compression hooks for plugin extensibility -还未自测，但是修改内容不多，如果可以的话可以直接合并 #5291

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
fbb7d27
feat: add context compression hooks for plugin extensibility
evpeople Feb 21, 2026
fcf482a
fix: improve context compression hooks based on code review
evpeople Feb 21, 2026
2ffe34b
fix: resolve circular import issue with lazy imports
evpeople Feb 21, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 46 additions & 3 deletions astrbot/core/agent/context/manager.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
from astrbot import logger
from astrbot.core.platform.astr_message_event import AstrMessageEvent

from ..message import Message
from .compressor import LLMSummaryCompressor, TruncateByTurnsCompressor
Expand Down Expand Up @@ -42,12 +43,17 @@ def __init__(
)

async def process(
self, messages: list[Message], trusted_token_usage: int = 0
self,
messages: list[Message],
trusted_token_usage: int = 0,
event: AstrMessageEvent | None = None,
) -> list[Message]:
"""Process the messages.

Args:
messages: The original message list.
trusted_token_usage: Trusted token usage from conversation.
event: Optional event for triggering hooks.

Returns:
The processed message list.
Expand All @@ -72,28 +78,47 @@ async def process(
if self.compressor.should_compress(
result, total_tokens, self.config.max_context_tokens
):
result = await self._run_compression(result, total_tokens)
result = await self._run_compression(result, total_tokens, event)

return result
except Exception as e:
Copy link
Contributor

@sourcery-ai sourcery-ai bot Feb 28, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

suggestion: 考虑在记录钩子失败时带上完整的 traceback，以便更容易调试。

使用 logger.warning(..., exc_info=e)（或直接使用 logger.exception(...)）可以保留 traceback，而不只是错误信息，这会让钩子失败问题更容易定位，同时仍然将其视作非致命错误。请在前置钩子和后置钩子中都采用这一做法。

建议实现：

            return result
        except Exception as e:
            logger.warning("Context compression hook failed", exc_info=e)
            return messages
  1. 在本文件中对应的“前置钩子”的 try/except 中也采用同样的模式，例如：
    • 将任何 except Exception as e: return ... 替换为：
      except Exception as e: logger.warning("Context before-hook failed", exc_info=e); return ...
      并使用合适的日志消息字符串。
  2. 如果该模块当前尚未定义 logger，请在文件顶部添加：
    • import logging
    • logger = logging.getLogger(__name__)
Original comment in English

suggestion: Consider logging hook failures with full traceback for easier debugging.

Using logger.warning(..., exc_info=e) (or just logger.exception(...)) will preserve the traceback instead of only the message, which makes hook failures much easier to diagnose while still treating them as non-fatal. Apply this to both the before and after hooks.

Suggested implementation:

            return result
        except Exception as e:
            logger.warning("Context compression hook failed", exc_info=e)
            return messages
  1. Apply the same pattern to the corresponding "before hook" try/except in this file, e.g.:
    • Replace any except Exception as e: return ... with:
      except Exception as e: logger.warning("Context before-hook failed", exc_info=e); return ...
      using an appropriate message string.
  2. If this module does not already define logger, add at the top of the file:
    • import logging
    • logger = logging.getLogger(__name__)

logger.error(f"Error during context processing: {e}", exc_info=True)
return messages

async def _run_compression(
self, messages: list[Message], prev_tokens: int
self,
messages: list[Message],
prev_tokens: int,
event: AstrMessageEvent | None = None,
) -> list[Message]:
"""
Compress/truncate the messages.

Args:
messages: The original message list.
prev_tokens: The token count before compression.
event: Optional event for triggering hooks.

Returns:
The compressed/truncated message list.
"""
logger.debug("Compress triggered, starting compression...")

# Trigger before compression hook
if event:
try:
from astrbot.core.pipeline.context_utils import call_event_hook
from astrbot.core.star.star_handler import EventType

await call_event_hook(
event,
EventType.OnBeforeContextCompressionEvent,
messages,
prev_tokens,
)
except Exception as e:
logger.warning(f"Hook OnBeforeContextCompressionEvent failed: {e}")

messages = await self.compressor(messages)

# double check
Expand All @@ -117,4 +142,22 @@ async def _run_compression(
# still need compress, truncate by half
messages = self.truncator.truncate_by_halving(messages)

# Recalculate token count after all truncation steps
final_tokens = self.token_counter.count_tokens(messages)

# Trigger after compression hook
if event:
try:
from astrbot.core.pipeline.context_utils import call_event_hook
from astrbot.core.star.star_handler import EventType

await call_event_hook(
event,
EventType.OnAfterContextCompressionEvent,
messages,
final_tokens,
)
except Exception as e:
logger.warning(f"Hook OnAfterContextCompressionEvent failed: {e}")

return messages
3 changes: 2 additions & 1 deletion astrbot/core/agent/runners/tool_loop_agent_runner.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -296,8 +296,9 @@ async def step(self):
# do truncate and compress
token_usage = self.req.conversation.token_usage if self.req.conversation else 0
self._simple_print_message_role("[BefCompact]")
event = getattr(self.run_context.context, "event", None)
self.run_context.messages = await self.context_manager.process(
self.run_context.messages, trusted_token_usage=token_usage
self.run_context.messages, trusted_token_usage=token_usage, event=event
)
self._simple_print_message_role("[AftCompact]")

Expand Down
2 changes: 2 additions & 0 deletions astrbot/core/star/star_handler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -201,6 +201,8 @@ class EventType(enum.Enum):
OnLLMToolRespondEvent = enum.auto() # 调用函数工具后
OnAfterMessageSentEvent = enum.auto() # 发送消息后
OnPluginErrorEvent = enum.auto() # 插件处理消息异常时
OnBeforeContextCompressionEvent = enum.auto() # 上下文压缩前事件
OnAfterContextCompressionEvent = enum.auto() # 上下文压缩后事件


H = TypeVar("H", bound=Callable[..., Any])
Expand Down