From fbb7d272a3a98acf7d16c6aad8ef7bfc1443df87 Mon Sep 17 00:00:00 2001 From: hangcaihui Date: Sat, 21 Feb 2026 21:00:51 +0800 Subject: [PATCH 1/3] feat: add context compression hooks for plugin extensibility Add two new event types: - OnBeforeContextCompressionEvent: triggered before context compression - OnAfterContextCompressionEvent: triggered after context compression Plugins can now register handlers to execute custom logic during context compression, enabling features like logging, analytics, or preserving specific conversation history. --- astrbot/core/agent/context/manager.py | 36 +++++++++++++++++-- .../agent/runners/tool_loop_agent_runner.py | 7 +++- astrbot/core/star/star_handler.py | 2 ++ 3 files changed, 41 insertions(+), 4 deletions(-) diff --git a/astrbot/core/agent/context/manager.py b/astrbot/core/agent/context/manager.py index 216a3e7e15..2e639f0a11 100644 --- a/astrbot/core/agent/context/manager.py +++ b/astrbot/core/agent/context/manager.py @@ -1,4 +1,7 @@ from astrbot import logger +from astrbot.core.pipeline.context_utils import call_event_hook +from astrbot.core.platform.astr_message_event import AstrMessageEvent +from astrbot.core.star.star_handler import EventType from ..message import Message from .compressor import LLMSummaryCompressor, TruncateByTurnsCompressor @@ -42,12 +45,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. @@ -72,7 +80,7 @@ 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: @@ -80,7 +88,10 @@ async def process( 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. @@ -88,12 +99,22 @@ async def _run_compression( 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: + await call_event_hook( + event, + EventType.OnBeforeContextCompressionEvent, + messages, + prev_tokens, + ) + messages = await self.compressor(messages) # double check @@ -117,4 +138,13 @@ async def _run_compression( # still need compress, truncate by half messages = self.truncator.truncate_by_halving(messages) + # Trigger after compression hook + if event: + await call_event_hook( + event, + EventType.OnAfterContextCompressionEvent, + messages, + tokens_after_summary, + ) + return messages diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 9f80dae1c9..7a241f2e52 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -296,8 +296,13 @@ 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 = ( + self.run_context.context.event + if hasattr(self.run_context.context, "event") + else 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]") diff --git a/astrbot/core/star/star_handler.py b/astrbot/core/star/star_handler.py index 63b0c447de..75ea448125 100644 --- a/astrbot/core/star/star_handler.py +++ b/astrbot/core/star/star_handler.py @@ -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]) From fcf482a9390ad7f113e90e6453ad0f21863711d7 Mon Sep 17 00:00:00 2001 From: hangcaihui Date: Sat, 21 Feb 2026 21:19:53 +0800 Subject: [PATCH 2/3] fix: improve context compression hooks based on code review - Wrap hook calls in try/except to prevent plugin failures from breaking core compression behavior - Recalculate token count after all truncation steps for accurate reporting - Simplify event extraction using getattr with default value --- astrbot/core/agent/context/manager.py | 33 ++++++++++++------- .../agent/runners/tool_loop_agent_runner.py | 6 +--- 2 files changed, 22 insertions(+), 17 deletions(-) diff --git a/astrbot/core/agent/context/manager.py b/astrbot/core/agent/context/manager.py index 2e639f0a11..b8add6ffe3 100644 --- a/astrbot/core/agent/context/manager.py +++ b/astrbot/core/agent/context/manager.py @@ -108,12 +108,15 @@ async def _run_compression( # Trigger before compression hook if event: - await call_event_hook( - event, - EventType.OnBeforeContextCompressionEvent, - messages, - prev_tokens, - ) + try: + 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) @@ -138,13 +141,19 @@ 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: - await call_event_hook( - event, - EventType.OnAfterContextCompressionEvent, - messages, - tokens_after_summary, - ) + try: + await call_event_hook( + event, + EventType.OnAfterContextCompressionEvent, + messages, + final_tokens, + ) + except Exception as e: + logger.warning(f"Hook OnAfterContextCompressionEvent failed: {e}") return messages diff --git a/astrbot/core/agent/runners/tool_loop_agent_runner.py b/astrbot/core/agent/runners/tool_loop_agent_runner.py index 7a241f2e52..1062e87a1e 100644 --- a/astrbot/core/agent/runners/tool_loop_agent_runner.py +++ b/astrbot/core/agent/runners/tool_loop_agent_runner.py @@ -296,11 +296,7 @@ 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 = ( - self.run_context.context.event - if hasattr(self.run_context.context, "event") - else None - ) + 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, event=event ) From 2ffe34b8efa3128fe61d5591084632cbcbfbd61e Mon Sep 17 00:00:00 2001 From: hangcaihui Date: Sat, 21 Feb 2026 21:27:38 +0800 Subject: [PATCH 3/3] fix: resolve circular import issue with lazy imports Move call_event_hook and EventType imports inside the hook functions to avoid circular import at module load time. --- astrbot/core/agent/context/manager.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/astrbot/core/agent/context/manager.py b/astrbot/core/agent/context/manager.py index b8add6ffe3..d95b99a7e7 100644 --- a/astrbot/core/agent/context/manager.py +++ b/astrbot/core/agent/context/manager.py @@ -1,7 +1,5 @@ from astrbot import logger -from astrbot.core.pipeline.context_utils import call_event_hook from astrbot.core.platform.astr_message_event import AstrMessageEvent -from astrbot.core.star.star_handler import EventType from ..message import Message from .compressor import LLMSummaryCompressor, TruncateByTurnsCompressor @@ -109,6 +107,9 @@ async def _run_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, @@ -147,6 +148,9 @@ async def _run_compression( # 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,