fix: bound conversation compaction budgets

deeptutor/services/session/context_builder.py

@@ -1,6 +1,4 @@
-"""
-Build bounded conversation history for unified chat sessions.
-"""
+"""Build budgeted conversation history for unified chat sessions."""
 
 from __future__ import annotations
 
@@ -11,14 +9,27 @@
 from deeptutor.core.stream import StreamEvent, StreamEventType
 from deeptutor.core.trace import build_trace_metadata, merge_trace_metadata, new_call_id
 from deeptutor.services.llm.config import LLMConfig
-from deeptutor.services.llm.context_window import resolve_effective_context_window
+from deeptutor.services.llm.context_window import (
+    coerce_positive_int,
+    resolve_effective_context_window,
+)
 
 from .protocol import SessionStoreProtocol
 
 #: When the summarizer's output lands within this fraction of its hard token
 #: cap, assume the provider cut it mid-sentence and trim the partial tail.
 TRUNCATION_GUARD_RATIO = 0.95
 
+# Ratio-only budgets become impractical for models with very large context
+# windows: a 1M-token window would otherwise reserve hundreds of thousands of
+# tokens for chat history and ask the summarizer for a six-figure response.
+# Keep the rolling-summary strategy, but bound its history plan, summary
+# output, and raw-rebuild eligibility threshold. The summary ceiling is also
+# constrained by the active generation limit when that setting is lower.
+MAX_HISTORY_PLAN_TOKENS = 131_072
+MAX_SUMMARY_OUTPUT_TOKENS = 16_384
+MAX_RAW_REBUILD_TOKENS = 131_072
+
 
 def count_tokens(text: str) -> int:
     """Estimate token count with tiktoken when available."""
@@ -102,7 +113,11 @@ async def process(self, *_args, **_kwargs) -> dict[str, Any]:
 
 
 class ContextBuilder:
-    """Construct a bounded conversation history plus optional summary trace."""
+    """Construct history against a bounded plan plus optional summary trace.
+
+    The budget is a planning target, not destructive truncation: the newest
+    non-empty message is retained even when that single message exceeds it.
+    """
 
     def __init__(
         self,
@@ -123,18 +138,29 @@ def _effective_context_window(self, llm_config: LLMConfig) -> int:
 
     def _history_budget(self, llm_config: LLMConfig) -> int:
         effective_context_window = self._effective_context_window(llm_config)
-        return max(256, int(effective_context_window * self.history_budget_ratio))
-
-    def _summary_budget(self, budget: int) -> int:
-        return max(96, int(budget * self.summary_target_ratio))
+        ratio_budget = max(256, int(effective_context_window * self.history_budget_ratio))
+        return min(ratio_budget, MAX_HISTORY_PLAN_TOKENS)
+
+    def _summary_budget(self, budget: int, llm_config: LLMConfig | None = None) -> int:
+        ratio_budget = max(96, int(budget * self.summary_target_ratio))
+        output_cap = MAX_SUMMARY_OUTPUT_TOKENS
+        if llm_config is not None:
+            generation_limit = coerce_positive_int(getattr(llm_config, "max_tokens", None))
+            if generation_limit is not None:
+                output_cap = min(output_cap, generation_limit)
+        return min(ratio_budget, output_cap)
 
     def _recent_budget(self, budget: int) -> int:
-        return max(128, budget - self._summary_budget(budget))
+        # Keep the original ratio-based split independent of the summarizer's
+        # output cap. Otherwise lowering that cap expands the verbatim tail and
+        # leaves almost no headroom before the next compaction.
+        return max(128, int(budget * (1 - self.summary_target_ratio)))
 
     def _rebuild_source_budget(self, llm_config: LLMConfig) -> int:
-        # Raw-rebuild input may use up to half the effective context window;
-        # beyond that we degrade to fold-in (existing summary + new turns).
-        return max(1024, self._effective_context_window(llm_config) // 2)
+        # A raw prefix is eligible for drift-free rebuild up to this threshold;
+        # beyond it we degrade to fold-in (existing summary + new turns).
+        ratio_budget = max(1024, self._effective_context_window(llm_config) // 2)
+        return min(ratio_budget, MAX_RAW_REBUILD_TOKENS)
 
     def _build_history(self, summary: str, messages: list[dict[str, Any]]) -> list[dict[str, Any]]:
         history: list[dict[str, Any]] = []
@@ -282,7 +308,7 @@ async def _trace_bridge(update: dict[str, Any]) -> None:
         )
         # The instruction targets ~80% of the hard cap so the model's own
         # length control — not the max_tokens cut — is the binding limit.
-        target_tokens = max(96, int(summary_budget * 0.8))
+        target_tokens = max(1, int(summary_budget * 0.8))
         system_prompt = (
             "You maintain a running summary of a conversation so future turns can "
             "continue seamlessly. Rewrite the summary from the material provided, "
@@ -365,7 +391,7 @@ async def build(
             return ContextBuildResult([], "", "", [], 0, self._history_budget(llm_config))
 
         budget = self._history_budget(llm_config)
-        summary_budget = self._summary_budget(budget)
+        summary_budget = self._summary_budget(budget, llm_config)
         recent_budget = self._recent_budget(budget)
 
         stored_summary = str(session.get("compressed_summary", "") or "").strip()

tests/services/session/test_context_builder.py

@@ -8,6 +8,9 @@
 import pytest
 
 from deeptutor.services.session.context_builder import (
+    MAX_HISTORY_PLAN_TOKENS,
+    MAX_RAW_REBUILD_TOKENS,
+    MAX_SUMMARY_OUTPUT_TOKENS,
     ContextBuilder,
     ContextBuildResult,
     build_history_text,
@@ -152,6 +155,11 @@ def test_history_budget_uses_large_context_model_heuristic(self) -> None:
         budget = builder._history_budget(self._make_llm_config(4096, model="gpt-4o-mini"))
         assert budget == int(65536 * 0.35)
 
+    def test_history_budget_has_absolute_cap(self) -> None:
+        builder = ContextBuilder(store=MagicMock(), history_budget_ratio=0.35)
+        budget = builder._history_budget(self._make_llm_config(4096, context_window=983_616))
+        assert budget == MAX_HISTORY_PLAN_TOKENS
+
     def test_history_budget_minimum(self) -> None:
         builder = ContextBuilder(store=MagicMock(), history_budget_ratio=0.01)
         budget = builder._history_budget(self._make_llm_config(100, model="unknown-local-model"))
@@ -165,6 +173,20 @@ def test_summary_budget_minimum(self) -> None:
         builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.01)
         assert builder._summary_budget(100) >= 96
 
+    def test_summary_budget_has_absolute_cap(self) -> None:
+        builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40)
+        assert builder._summary_budget(MAX_HISTORY_PLAN_TOKENS) == (MAX_SUMMARY_OUTPUT_TOKENS)
+
+    def test_summary_budget_respects_configured_output_limit(self) -> None:
+        builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40)
+        cfg = self._make_llm_config(4096, context_window=983_616)
+        assert builder._summary_budget(MAX_HISTORY_PLAN_TOKENS, cfg) == 4096
+
+    def test_summary_budget_can_follow_small_generation_limit(self) -> None:
+        builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40)
+        cfg = self._make_llm_config(50, context_window=983_616)
+        assert builder._summary_budget(MAX_HISTORY_PLAN_TOKENS, cfg) == 50
+
     def test_recent_budget(self) -> None:
         builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40)
         recent = builder._recent_budget(1000)
@@ -174,6 +196,20 @@ def test_recent_budget_minimum(self) -> None:
         builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.99)
         assert builder._recent_budget(200) >= 128
 
+    def test_recent_budget_preserves_headroom_when_summary_is_capped(self) -> None:
+        builder = ContextBuilder(store=MagicMock(), summary_target_ratio=0.40)
+        cfg = self._make_llm_config(4096, context_window=983_616)
+        recent_budget = builder._recent_budget(MAX_HISTORY_PLAN_TOKENS)
+        summary_budget = builder._summary_budget(MAX_HISTORY_PLAN_TOKENS, cfg)
+
+        assert recent_budget == int(MAX_HISTORY_PLAN_TOKENS * 0.60)
+        assert recent_budget + summary_budget < MAX_HISTORY_PLAN_TOKENS
+
+    def test_rebuild_source_budget_has_absolute_cap(self) -> None:
+        builder = ContextBuilder(store=MagicMock())
+        cfg = self._make_llm_config(4096, context_window=983_616)
+        assert builder._rebuild_source_budget(cfg) == MAX_RAW_REBUILD_TOKENS
+
 
 # ---------------------------------------------------------------------------
 # ContextBuilder._build_history
@@ -328,6 +364,50 @@ async def test_large_context_model_avoids_premature_summarize(self) -> None:
         assert result.events == []
         store.update_summary.assert_not_called()
 
+    @pytest.mark.asyncio
+    async def test_large_context_build_uses_capped_plan_with_headroom(self) -> None:
+        store = AsyncMock()
+        store.get_session = AsyncMock(
+            return_value={
+                "id": "s1",
+                "compressed_summary": "",
+                "summary_up_to_msg_id": 0,
+            }
+        )
+        store.get_messages_for_context = AsyncMock(
+            return_value=[
+                {"id": 1, "role": "user", "content": "old turn"},
+                {"id": 2, "role": "assistant", "content": "recent turn"},
+            ]
+        )
+
+        builder = ContextBuilder(store=store)
+        builder._summarize = AsyncMock(return_value=("SUMMARY", []))
+        select_recent = MagicMock(wraps=builder._select_recent_messages)
+        builder._select_recent_messages = select_recent
+        cfg = MagicMock()
+        cfg.max_tokens = 4096
+        cfg.model = "qwen3.8-max"
+        cfg.context_window = 983_616
+
+        def _fake_count_tokens(text: str) -> int:
+            if text in {"old turn", "recent turn"}:
+                return 70_000
+            if "User: old turn" in text and "Assistant: recent turn" in text:
+                return 140_000
+            return 100
+
+        with patch(
+            "deeptutor.services.session.context_builder.count_tokens",
+            side_effect=_fake_count_tokens,
+        ):
+            result = await builder.build(session_id="s1", llm_config=cfg)
+
+        assert result.budget == MAX_HISTORY_PLAN_TOKENS
+        assert builder._summarize.call_args.kwargs["summary_budget"] == 4096
+        assert select_recent.call_args.args[1] == int(MAX_HISTORY_PLAN_TOKENS * 0.60)
+        store.update_summary.assert_awaited_once_with("s1", "SUMMARY", 1)
+
 
 # ---------------------------------------------------------------------------
 # trim_incomplete_tail
@@ -342,6 +422,40 @@ def test_single_line_kept(self) -> None:
         assert trim_incomplete_tail("only one line, keep it") == "only one line, keep it"
 
 
+# ---------------------------------------------------------------------------
+# ContextBuilder._summarize
+# ---------------------------------------------------------------------------
+
+
+class TestContextBuilderSummarize:
+    @pytest.mark.asyncio
+    async def test_target_stays_below_small_output_cap(self) -> None:
+        captured: dict[str, Any] = {}
+
+        async def _stream_llm(**kwargs: Any):
+            captured.update(kwargs)
+            yield "short summary"
+
+        agent = MagicMock()
+        agent.stream_llm = _stream_llm
+        builder = ContextBuilder(store=MagicMock())
+
+        with patch(
+            "deeptutor.services.session.context_builder._ContextSummaryAgent",
+            return_value=agent,
+        ):
+            summary, _events = await builder._summarize(
+                session_id="s1",
+                language="en",
+                source_text="User: hello",
+                summary_budget=50,
+            )
+
+        assert summary == "short summary"
+        assert captured["max_tokens"] == 50
+        assert "under 40 tokens" in captured["user_prompt"]
+
+
 # ---------------------------------------------------------------------------
 # ContextBuilder.build — summarize paths (rebuild / fold-in / failure / branch)
 # ---------------------------------------------------------------------------