fix: fall back from a failed context summary without long retries (#10436)

* fix: fall back from a failed context summary without long retries

The summary request in LLMSummaryCompressor used the provider's default
retries, so an exhausted quota (HTTP 429) was retried five times with
backoff before compression fell back to truncation, holding up the reply
even after switching the chat model.

Request the summary once, since a failure already falls back to
truncation. When the context still exceeds the limit, drop whole rounds,
oldest first, and recount after each one until it fits, instead of
halving the messages once without checking again. The latest round is
always kept.

Fixes #10433

* fix: leave the token count out of the fallback warning

CodeQL flags values from count_tokens() as sensitive data, so logging the
count in the new warning raised a clear-text logging alert. The count is
already in the "Compress completed" line just before, so drop it here.

* fix: log how many rounds the context fallback dropped
This commit is contained in:
Niansia
2026-10-08 22:43:14 +09:00
committed by GitHub
parent 30820a48fa
commit a8bb0803ef
3 changed files with 156 additions and 27 deletions
+4 -1
View File
@@ -271,10 +271,13 @@ class LLMSummaryCompressor:
)
log_context_sanitize_stats(sanitize_stats)
# Generate summary
# Generate summary. A single attempt is enough: on failure the caller
# falls back to truncation, so retrying (e.g. on an exhausted quota)
# would only hold up the reply.
try:
response = await self.provider.text_chat(
contexts=sanitized_summary_contexts,
request_max_retries=1,
)
summary_content = (response.completion_text or "").strip()
except Exception as e:
+41 -6
View File
@@ -3,6 +3,7 @@ from astrbot import logger
from ..message import Message
from .compressor import LLMSummaryCompressor, TruncateByTurnsCompressor
from .config import ContextConfig
from .round_utils import split_into_rounds
from .token_counter import EstimateTokenCounter
from .truncator import ContextTruncator
@@ -86,6 +87,11 @@ class ContextManager:
"""
Compress/truncate the messages.
If the context still exceeds the limit after compression (for example
because the summary request failed), whole rounds are dropped, oldest
first, until it fits. The latest round is always kept, even when it
alone exceeds the limit.
Args:
messages: The original message list.
prev_tokens: The token count before compression.
@@ -109,13 +115,42 @@ class ContextManager:
)
# last check
if self.compressor.should_compress(
if not self.compressor.should_compress(
messages, tokens_after_summary, self.config.max_context_tokens
):
logger.info(
"Context still exceeds max tokens after compression, applying halving truncation..."
)
# still need compress, truncate by half
messages = self.truncator.truncate_by_halving(messages)
return messages
logger.info(
"Context still exceeds max tokens after compression, dropping the oldest rounds..."
)
# Drop whole rounds, oldest first, so that tool calls stay with their
# results, and recount after each one. The latest round is the current
# request and is always kept.
first_non_system = next(
(i for i, msg in enumerate(messages) if msg.role != "system"),
len(messages),
)
system_messages = messages[:first_non_system]
rounds = [
[seg for seg in rnd if isinstance(seg, Message)]
for rnd in split_into_rounds(messages[first_non_system:])
]
tokens = tokens_after_summary
dropped_rounds = 0
while len(rounds) > 1 and self.compressor.should_compress(
messages, tokens, self.config.max_context_tokens
):
rounds.pop(0)
dropped_rounds += 1
messages = system_messages + [msg for rnd in rounds for msg in rnd]
tokens = self.token_counter.count_tokens(messages)
logger.info(f"Dropped {dropped_rounds} oldest round(s).")
if self.compressor.should_compress(
messages, tokens, self.config.max_context_tokens
):
logger.warning(
"Context still exceeds max tokens with only the latest round left; "
"sending it as is."
)
return messages
+111 -20
View File
@@ -12,8 +12,15 @@ sys.path.insert(0, str(Path(__file__).parent.parent.parent))
from astrbot.core.agent.context.config import ContextConfig
from astrbot.core.agent.context.manager import ContextManager
from astrbot.core.agent.message import AudioURLPart, ImageURLPart, Message, TextPart
from astrbot.core.agent.message import (
AudioURLPart,
ImageURLPart,
Message,
TextPart,
ToolCall,
)
from astrbot.core.provider.entities import LLMResponse
from astrbot.core.provider.sources.request_retry import retry_provider_request
class MockProvider:
@@ -42,6 +49,44 @@ class MockProvider:
return MagicMock(id="test_provider", type="openai")
class FailingProvider(MockProvider):
"""Provider whose summary request always fails."""
async def text_chat(self, **kwargs):
self.last_text_chat_kwargs = kwargs
raise RuntimeError("summary model unavailable")
class QuotaExhaustedError(Exception):
status_code = 429
class RateLimitedProvider(MockProvider):
"""Provider that goes through the real retry helper and always gets a 429."""
def __init__(self):
super().__init__()
self.attempts = 0
async def text_chat(self, **kwargs):
async def request():
self.attempts += 1
raise QuotaExhaustedError("quota exhausted")
return await retry_provider_request(
"summary",
request,
max_attempts=kwargs.get("request_max_retries"),
)
class CharTokenCounter:
"""Counts one token per character of text content."""
def count_tokens(self, messages, trusted_token_usage: int = 0) -> int:
return sum(len(m.content) for m in messages if isinstance(m.content, str))
class TestContextManager:
"""Test suite for ContextManager."""
@@ -580,31 +625,77 @@ class TestContextManager:
mock_compress.assert_not_called()
assert result == messages
def _rounds_with_tool_chain(self, latest_request: str) -> list[Message]:
"""System prompt, a long first round, a round with a tool chain, and
the latest request. One token per character with ``CharTokenCounter``."""
return [
self.create_message("system", "s" * 10),
self.create_message("user", "a" * 300),
self.create_message("assistant", "b" * 300),
self.create_message("user", "c" * 50),
Message(
role="assistant",
tool_calls=[
ToolCall(
id="call_1",
function=ToolCall.FunctionBody(name="lookup", arguments="{}"),
)
],
),
Message(role="tool", content="d" * 100, tool_call_id="call_1"),
self.create_message("assistant", "e" * 50),
self.create_message("user", latest_request),
]
@pytest.mark.asyncio
async def test_double_check_after_compression(self):
"""Test that halving is applied if still over threshold after compression."""
config = ContextConfig(max_context_tokens=100)
async def test_failed_summary_drops_oldest_rounds_until_within_limit(self):
"""When the summary fails, whole rounds are dropped until the context fits."""
provider = FailingProvider()
config = ContextConfig(
max_context_tokens=1000,
llm_compress_provider=provider, # type: ignore[arg-type]
custom_token_counter=CharTokenCounter(),
)
manager = ContextManager(config)
messages = self._rounds_with_tool_chain("f" * 20)
# Create messages that would still be over threshold after compression
long_messages = [self.create_message("user", "x" * 200) for _ in range(10)]
result = await manager.process(messages)
# Mock compressor to return messages still over threshold
async def mock_compress(msgs):
return msgs # Return same messages (still over limit)
assert provider.last_text_chat_kwargs["request_max_retries"] == 1
# Dropping the first round (600 tokens) is enough to get under 820, so
# the round with the tool chain stays intact.
assert result == [messages[0], *messages[3:]]
# Mock should_compress to return True twice (before and after compression)
with patch.object(manager.compressor, "should_compress", return_value=True):
with patch.object(manager.compressor, "__call__", new=mock_compress):
with patch.object(
manager.truncator,
"truncate_by_halving",
return_value=long_messages[:5],
) as mock_halving:
_ = await manager.process(long_messages)
@pytest.mark.asyncio
async def test_fallback_keeps_latest_round_when_it_alone_exceeds_limit(
self, caplog
):
config = ContextConfig(
max_context_tokens=1000,
llm_compress_provider=FailingProvider(), # type: ignore[arg-type]
custom_token_counter=CharTokenCounter(),
)
manager = ContextManager(config)
messages = self._rounds_with_tool_chain("f" * 900)
# Halving should be called
mock_halving.assert_called_once()
result = await manager.process(messages)
assert result == [messages[0], messages[-1]]
assert "only the latest round left" in caplog.text
@pytest.mark.asyncio
async def test_llm_summary_is_requested_once(self):
"""A failed summary falls back to truncation instead of being retried."""
from astrbot.core.agent.context.compressor import LLMSummaryCompressor
provider = RateLimitedProvider()
compressor = LLMSummaryCompressor(provider=provider, keep_recent_ratio=0.15) # type: ignore[arg-type]
messages = self._rounds_with_tool_chain("f" * 20)
result = await compressor(messages)
assert result is messages
assert provider.attempts == 1
# ==================== Combined Truncation and Compression Tests ====================