From 4bfe6c441be8b401fa3e3e441b23c686ff1f5b9e Mon Sep 17 00:00:00 2001 From: Sam Schillace Date: Thu, 6 Aug 2026 15:44:28 -0700 Subject: [PATCH 1/4] feat: add persistent context compaction --- amplifier_module_context_simple/__init__.py | 158 +++++++++++++++++--- tests/test_persistent_compaction.py | 131 ++++++++++++++++ 2 files changed, 272 insertions(+), 17 deletions(-) create mode 100644 tests/test_persistent_compaction.py diff --git a/amplifier_module_context_simple/__init__.py b/amplifier_module_context_simple/__init__.py index ddfdad4..e3a49c5 100644 --- a/amplifier_module_context_simple/__init__.py +++ b/amplifier_module_context_simple/__init__.py @@ -29,6 +29,20 @@ logger = logging.getLogger(__name__) +# Compaction observability events. Prefer the canonical constants from +# amplifier_core.events; fall back to the string literals if unavailable so the +# module keeps working standalone (e.g. in isolated unit tests). +try: # pragma: no cover - trivial import guard + from amplifier_core.events import ( + CONTEXT_COMPACTION as _EVT_COMPACTION, + CONTEXT_POST_COMPACT as _EVT_POST_COMPACT, + CONTEXT_PRE_COMPACT as _EVT_PRE_COMPACT, + ) +except Exception: # pragma: no cover + _EVT_PRE_COMPACT = "context:pre_compact" + _EVT_POST_COMPACT = "context:post_compact" + _EVT_COMPACTION = "context:compaction" + async def mount(coordinator: ModuleCoordinator, config: dict[str, Any] | None = None): """ @@ -346,23 +360,137 @@ async def clear(self) -> None: self.messages = [] logger.info("Context cleared") - async def should_compact(self) -> bool: - """Check if context should be compacted. + async def should_compact( + self, + provider: Any | None = None, + token_budget: int | None = None, + ) -> bool: + """Check whether the stored context is over the compaction threshold. + + Unlike ``get_messages_for_request`` (which compacts ephemerally every + request without touching stored state), this reports on the *persistent* + history in ``self.messages`` - i.e. whether an explicit/persistent + ``compact()`` would have anything to do. + + Args: + provider: Optional provider for dynamic budget calculation. + token_budget: Optional explicit budget override. + + Returns: + True if stored context usage >= compact_threshold. + """ + report = self.usage_report(provider=provider, token_budget=token_budget) + return report["tokens"] >= report["threshold_tokens"] + + def usage_report( + self, + provider: Any | None = None, + token_budget: int | None = None, + ) -> dict[str, Any]: + """Return a snapshot of current context token usage vs budget. - Note: This module uses ephemeral compaction during get_messages_for_request(), - so this always returns False. The actual compaction check happens internally. - This method exists to satisfy the ContextManager protocol. + Intended for status displays and for deciding whether to compact. + Uses the same heuristic estimator and budget calculation as compaction. """ - return False + budget = self._calculate_budget(token_budget, provider) + tokens = self._estimate_tokens(list(self.messages)) + pct = (tokens / budget) if budget > 0 else 0.0 + return { + "tokens": tokens, + "budget": budget, + "pct": pct, + "threshold": self.compact_threshold, + "threshold_tokens": int(budget * self.compact_threshold), + "target_tokens": int(budget * self.target_usage), + "messages": len(self.messages), + } + + async def compact( + self, + provider: Any | None = None, + token_budget: int | None = None, + force: bool = False, + ) -> dict[str, Any]: + """Persistently compact the stored conversation history. - async def compact(self) -> None: - """Compact the context. + This is distinct from the ephemeral compaction applied inside + ``get_messages_for_request()``: it runs the same progressive strategy + but COMMITS the result back to ``self.messages``, permanently shrinking + the stored history (and therefore future transcripts and every + subsequent request). System messages are always preserved. + + Args: + provider: Optional provider for dynamic budget calculation. + token_budget: Optional explicit budget override. + force: When True, compact even if below the trigger threshold + (used by the explicit ``/compact`` command). Even when forced, + nothing is removed if the history is already at/under the target. - Note: This module uses ephemeral compaction during get_messages_for_request(), - so this is a no-op. Compaction happens automatically when getting messages. - This method exists to satisfy the ContextManager protocol. + Returns: + A stats dict. ``compacted`` is False (with a ``reason``) when no + work was done; otherwise it carries before/after token and message + counts plus the strategy level reached. """ - pass + budget = self._calculate_budget(token_budget, provider) + working = list(self.messages) + before_tokens = self._estimate_tokens(working) + before_messages = len(working) + target_tokens = int(budget * self.target_usage) + + # Nothing to do: below threshold (unless forced) or already at target. + if not force and not self._should_compact(before_tokens, budget): + return { + "compacted": False, + "reason": "below_threshold", + "before_tokens": before_tokens, + "after_tokens": before_tokens, + "before_messages": before_messages, + "after_messages": before_messages, + "budget": budget, + } + if before_tokens <= target_tokens: + return { + "compacted": False, + "reason": "already_compact", + "before_tokens": before_tokens, + "after_tokens": before_tokens, + "before_messages": before_messages, + "after_messages": before_messages, + "budget": budget, + "target_tokens": target_tokens, + } + + # Emit pre-compact (best effort). + await self._emit(_EVT_PRE_COMPACT, {"before_tokens": before_tokens, "budget": budget}) + + # Reuse the progressive strategy, then commit the result. + compacted = await self._compact_ephemeral(budget, working) + self.messages = compacted + + stats = dict(self._last_compaction_stats or {}) + stats["compacted"] = True + stats.setdefault("before_tokens", before_tokens) + stats.setdefault("after_tokens", self._estimate_tokens(compacted)) + stats.setdefault("before_messages", before_messages) + stats.setdefault("after_messages", len(compacted)) + stats["persistent"] = True + stats["forced"] = force + + await self._emit(_EVT_POST_COMPACT, stats) + logger.info( + f"Persistent compaction committed: {before_messages} -> {len(compacted)} messages, " + f"{before_tokens:,} -> {stats['after_tokens']:,} tokens" + ) + return stats + + async def _emit(self, event: str, payload: dict[str, Any]) -> None: + """Emit a hook event if a hooks bus is available (best effort).""" + if self._hooks is None: + return + try: + await self._hooks.emit(event, payload) + except Exception as e: # pragma: no cover - observability must never break flow + logger.warning(f"Could not emit {event} event: {e}") def _should_compact(self, token_count: int, budget: int) -> bool: """Check if context should be compacted.""" @@ -1086,11 +1214,7 @@ async def _finalize_compaction_with_stats( self._last_compaction_stats = stats # Emit event if hooks available - if self._hooks is not None: - try: - await self._hooks.emit("context:compaction", stats) - except Exception as e: - logger.warning(f"Could not emit compaction event: {e}") + await self._emit(_EVT_COMPACTION, stats) return final_messages diff --git a/tests/test_persistent_compaction.py b/tests/test_persistent_compaction.py new file mode 100644 index 0000000..a2c8df2 --- /dev/null +++ b/tests/test_persistent_compaction.py @@ -0,0 +1,131 @@ +""" +Tests for the persistent compaction primitives: + +- should_compact(): reports on the STORED history (not ephemeral view) +- compact(force=...): COMMITS a compacted history back to self.messages +- usage_report(): snapshot of token usage vs budget +- pre/post compaction events are emitted on the hooks bus + +These are distinct from the ephemeral compaction applied during +get_messages_for_request(), which never mutates self.messages. +""" + +import pytest +from amplifier_module_context_simple import SimpleContextManager + + +class _RecordingHooks: + """Minimal hooks bus that records emitted events.""" + + def __init__(self): + self.events: list[tuple[str, dict]] = [] + + async def emit(self, event: str, payload: dict) -> None: + self.events.append((event, payload)) + + +async def _fill(context: SimpleContextManager, pairs: int) -> None: + for i in range(pairs): + await context.add_message( + {"role": "user", "content": f"message {i} with some padding content here"} + ) + await context.add_message( + {"role": "assistant", "content": f"response {i} with some padding content here"} + ) + + +@pytest.mark.asyncio +async def test_should_compact_reflects_stored_history(): + context = SimpleContextManager(max_tokens=1000, compact_threshold=0.9, target_usage=0.5) + assert await context.should_compact() is False + await _fill(context, 50) + assert await context.should_compact() is True + + +@pytest.mark.asyncio +async def test_compact_persists_and_shrinks_history(): + hooks = _RecordingHooks() + context = SimpleContextManager( + max_tokens=1000, + compact_threshold=0.9, + target_usage=0.5, + protected_recent=0.1, + hooks=hooks, + ) + await _fill(context, 50) + + before = len(context.messages) + before_tokens = context._estimate_tokens(context.messages) + + stats = await context.compact() + + assert stats["compacted"] is True + assert stats["persistent"] is True + # Persistent: self.messages is actually mutated (unlike ephemeral path). + assert len(context.messages) < before + assert context._estimate_tokens(context.messages) < before_tokens + assert stats["after_tokens"] <= stats["target_tokens"] or stats["strategy_level"] >= 1 + + # Pre/post events emitted. + names = [e[0] for e in hooks.events] + assert "context:pre_compact" in names + assert "context:post_compact" in names + + +@pytest.mark.asyncio +async def test_compact_noop_below_threshold(): + context = SimpleContextManager(max_tokens=100_000, compact_threshold=0.9, target_usage=0.5) + await _fill(context, 3) + before = len(context.messages) + + stats = await context.compact() # not forced, well below threshold + + assert stats["compacted"] is False + assert stats["reason"] == "below_threshold" + assert len(context.messages) == before # untouched + + +@pytest.mark.asyncio +async def test_force_compact_reports_already_compact_when_under_target(): + context = SimpleContextManager(max_tokens=100_000, compact_threshold=0.9, target_usage=0.5) + await _fill(context, 3) + before = len(context.messages) + + stats = await context.compact(force=True) + + # Forced past the threshold gate, but history is already under the target, + # so nothing is removed. + assert stats["compacted"] is False + assert stats["reason"] == "already_compact" + assert len(context.messages) == before + + +@pytest.mark.asyncio +async def test_system_messages_preserved_across_persistent_compaction(): + context = SimpleContextManager( + max_tokens=1000, compact_threshold=0.9, target_usage=0.5, protected_recent=0.1 + ) + await context.add_message( + {"role": "system", "content": "IDENTITY", "metadata": {"source": "hook"}} + ) + await _fill(context, 50) + + await context.compact() + + system_msgs = [m for m in context.messages if m.get("role") == "system"] + assert any(m.get("content") == "IDENTITY" for m in system_msgs), ( + "System/identity messages must survive persistent compaction" + ) + + +@pytest.mark.asyncio +async def test_usage_report_shape(): + context = SimpleContextManager(max_tokens=1000, compact_threshold=0.9, target_usage=0.5) + await _fill(context, 10) + report = context.usage_report() + assert report["budget"] == 1000 + assert report["threshold_tokens"] == 900 + assert report["target_tokens"] == 500 + assert report["tokens"] > 0 + assert 0.0 <= report["pct"] + assert report["messages"] == len(context.messages) From bed31ce3387449ebab76bd0d3542d5781182c42f Mon Sep 17 00:00:00 2001 From: Sam Schillace Date: Thu, 6 Aug 2026 16:00:20 -0700 Subject: [PATCH 2/4] fix: preserve compaction invariants --- README.md | 2 +- amplifier_module_context_simple/__init__.py | 82 +++++++++++++-------- 2 files changed, 52 insertions(+), 32 deletions(-) diff --git a/README.md b/README.md index 241acd2..a756e88 100644 --- a/README.md +++ b/README.md @@ -66,7 +66,7 @@ Not suitable for: ## Compaction Strategy -The SimpleContextManager uses **ephemeral compaction** - `get_messages_for_request()` returns a compacted VIEW without modifying the internal message history. The full history is always preserved in memory. +The SimpleContextManager uses **ephemeral request-time compaction**: `get_messages_for_request()` returns a compacted view without modifying the internal message history. Call `compact()` explicitly to persistently replace the stored history with a compacted version; subsequent transcripts then reflect that reduced history. Compaction triggers when token usage reaches the configured threshold (default: 92% of max_tokens): diff --git a/amplifier_module_context_simple/__init__.py b/amplifier_module_context_simple/__init__.py index e3a49c5..6c9b46a 100644 --- a/amplifier_module_context_simple/__init__.py +++ b/amplifier_module_context_simple/__init__.py @@ -1,13 +1,14 @@ """ Simple context manager module. -Implements an in-memory context manager with EPHEMERAL compaction: - • Messages stored in memory (self.messages is the source of truth) - • Compaction NEVER modifies self.messages - • get_messages_for_request() returns compacted VIEW (new list) - • get_messages() returns FULL history (for transcripts/session persistence) - -This design ensures conversation history is never lost, even during compaction. +Implements an in-memory context manager with two explicit compaction modes: + • Messages are stored in memory (self.messages is the source of truth) + • get_messages_for_request() returns an ephemeral compacted VIEW (new list) + • compact() explicitly persists compaction back to self.messages + • get_messages() returns the currently stored history + +Request-time compaction never changes stored history. Callers that invoke compact() +opt in to permanently shrinking the stored history and future transcripts. For persistent storage across sessions, use context-persistent instead. Dynamic System Prompt Support: @@ -90,15 +91,16 @@ async def mount(coordinator: ModuleCoordinator, config: dict[str, Any] | None = class SimpleContextManager: """ - In-memory context manager with EPHEMERAL compaction. + In-memory context manager with ephemeral and explicit persistent compaction. - Key Principle: self.messages is the source of truth and is NEVER modified - by compaction. Compaction only returns a compacted VIEW for the current - LLM request. + ``self.messages`` is the source of truth. Automatic request-time compaction + only returns a compacted view and leaves it unchanged; ``compact()`` is the + explicit operation that commits a compacted result back to stored history. Owns memory policy: orchestrators ask for messages via get_messages_for_request(), - and this context manager decides how to fit them within limits. Compaction is - handled internally and ephemerally - the original history is always preserved. + and this context manager decides how to fit them within limits. Normal request + preparation is ephemeral; persistent compaction happens only when a caller + explicitly invokes ``compact()``. Compaction Strategy (Progressive Interleaved): Triggered when usage >= compact_threshold (default 92%), target is target_usage (default 50%). @@ -239,11 +241,11 @@ async def get_messages_for_request( bundle instructions to be re-processed each turn. Applies EPHEMERAL compaction if needed - returns a NEW list without - modifying self.messages. The original history is always preserved. + modifying self.messages. If compaction occurs and notice is enabled, a system-reminder is inserted - at position 1 (after main system message) to inform the LLM about what - was compacted. + after any leading system messages and before the conversation. This avoids + separating an assistant tool call from its immediately following results. Args: token_budget: Optional explicit token limit (deprecated, prefer provider). @@ -320,9 +322,18 @@ async def get_messages_for_request( if level >= self.compaction_notice_min_level: notice = self._format_compaction_notice() if notice: - # Insert at position 1 (after main system message at position 0) + # Insert after leading system messages, before the first + # conversation message. A fixed index can split an + # assistant tool-call message from its following result + # when there is no leading system message. + notice_index = 0 + while ( + notice_index < len(compacted) + and compacted[notice_index].get("role") == "system" + ): + notice_index += 1 compacted.insert( - 1, + notice_index, { "role": "system", "content": notice, @@ -333,7 +344,7 @@ async def get_messages_for_request( }, ) logger.debug( - f"Inserted compaction notice at position 1 (level {level}, " + f"Inserted compaction notice at position {notice_index} (level {level}, " f"verbosity: {self.compaction_notice_verbosity})" ) @@ -343,10 +354,10 @@ async def get_messages_for_request( async def get_messages(self) -> list[dict[str, Any]]: """ - Get ALL messages (full history, never compacted) for transcripts/debugging. + Get the currently stored messages for transcripts/debugging. - This returns the complete, unmodified history - suitable for saving - to transcript files for session persistence. + Request-time compaction does not affect this list. An earlier explicit + ``compact()`` call may have persistently reduced it. """ return list(self.messages) @@ -463,19 +474,28 @@ async def compact( # Emit pre-compact (best effort). await self._emit(_EVT_PRE_COMPACT, {"before_tokens": before_tokens, "budget": budget}) - # Reuse the progressive strategy, then commit the result. + # Reuse the progressive strategy, then commit only an actual reduction. compacted = await self._compact_ephemeral(budget, working) - self.messages = compacted - stats = dict(self._last_compaction_stats or {}) - stats["compacted"] = True - stats.setdefault("before_tokens", before_tokens) - stats.setdefault("after_tokens", self._estimate_tokens(compacted)) - stats.setdefault("before_messages", before_messages) - stats.setdefault("after_messages", len(compacted)) + after_tokens = self._estimate_tokens(compacted) + after_messages = len(compacted) + stats["before_tokens"] = before_tokens + stats["after_tokens"] = after_tokens + stats["before_messages"] = before_messages + stats["after_messages"] = after_messages stats["persistent"] = True stats["forced"] = force + if after_tokens >= before_tokens and after_messages >= before_messages: + stats["compacted"] = False + stats["reason"] = "no_reduction" + logger.info( + "Persistent compaction made no reduction; stored history left unchanged" + ) + return stats + + self.messages = compacted + stats["compacted"] = True await self._emit(_EVT_POST_COMPACT, stats) logger.info( f"Persistent compaction committed: {before_messages} -> {len(compacted)} messages, " @@ -1316,7 +1336,7 @@ def _format_compaction_notice(self) -> str: What may be affected: {affected} -Note: This compaction is ephemeral (affects only this request). Full history is preserved in session transcript. +Note: This compaction is ephemeral (affects only this request). Stored history is unchanged by this request. """ return notice From 6ca1b572fbf1ac595202f1b5b394e99a15e09b42 Mon Sep 17 00:00:00 2001 From: Sam Schillace Date: Thu, 6 Aug 2026 16:32:50 -0700 Subject: [PATCH 3/4] fix: enforce persistent compaction boundaries --- amplifier_module_context_simple/__init__.py | 174 ++++++++++++-------- 1 file changed, 109 insertions(+), 65 deletions(-) diff --git a/amplifier_module_context_simple/__init__.py b/amplifier_module_context_simple/__init__.py index 6c9b46a..1ec4be4 100644 --- a/amplifier_module_context_simple/__init__.py +++ b/amplifier_module_context_simple/__init__.py @@ -24,6 +24,7 @@ import logging from collections.abc import Awaitable, Callable from datetime import UTC, datetime +from math import ceil from typing import Any from amplifier_core import ModuleCoordinator @@ -56,7 +57,7 @@ async def mount(coordinator: ModuleCoordinator, config: dict[str, Any] | None = - compact_threshold: Trigger compaction at this usage (default: 0.92) - target_usage: Compact down to this usage (default: 0.50) - protected_recent: Always protect last N% of messages (default: 0.30) - - protected_tool_results: Always protect last N tool results (default: 5) + - protected_tool_results: Protect last N recent tool results (default: 5) - truncate_chars: Characters to keep when truncating tool results (default: 250) - compaction_notice_enabled: Enable compaction notice (default: True) - compaction_notice_token_reserve: Tokens to reserve for notice (default: 800) @@ -120,7 +121,7 @@ class SimpleContextManager: - Preferring truncation (preserves structure) over removal (loses context) - Progressively relaxing protection as pressure increases - Respecting configured protected_recent as baseline, only relaxing under pressure - - Always protecting: system messages, last user message, last N tool results, tool pairs + - Always protecting: system messages, last user message, recent tool results, valid tool pairs - First user message: stubbable at Level 8, but never fully removed """ @@ -147,7 +148,7 @@ def __init__( compact_threshold: Trigger compaction at this usage ratio (0.0-1.0) target_usage: Compact down to this usage ratio (0.0-1.0) protected_recent: Always protect last N% of messages (0.0-1.0) - protected_tool_results: Always protect last N tool results from truncation + protected_tool_results: Protect last N results in the recent-history zone truncate_chars: Characters to keep when truncating tool results compaction_notice_enabled: Enable compaction notice injection compaction_notice_token_reserve: Tokens to reserve for notice @@ -255,30 +256,21 @@ async def get_messages_for_request( Returns: Messages ready for LLM request, compacted if necessary. """ + # The caller/configured value remains the governing request budget. + # A valid notice reserve may lower the conversation-compaction budget, + # but can never make it non-positive or replace the final request cap. budget = self._calculate_budget(token_budget, provider) - - # Reserve token budget for potential compaction notice (if enabled) - effective_budget = budget - if self.compaction_notice_enabled: - effective_budget = budget - self.compaction_notice_token_reserve - if effective_budget <= 0: - # Misconfiguration guard: if the reserve consumes the entire budget - # (or more), _should_compact's `budget > 0` check would silently - # force usage to 0, disabling compaction entirely rather than - # loudly failing. Fall back to the full budget instead of a - # non-positive effective budget - a reserve that swallows the - # whole context is not a valid state to compact against. - logger.warning( - f"compaction_notice_token_reserve ({self.compaction_notice_token_reserve:,}) " - f">= budget ({budget:,}); ignoring reserve for this request to avoid " - f"silently disabling compaction (effective budget would be {effective_budget:,})" - ) - effective_budget = budget - else: - logger.debug( - f"Reserved {self.compaction_notice_token_reserve} tokens for potential notice " - f"(effective budget: {effective_budget:,})" - ) + compaction_budget = budget + if ( + self.compaction_notice_enabled + and 0 < self.compaction_notice_token_reserve < budget + ): + compaction_budget = budget - self.compaction_notice_token_reserve + logger.debug( + f"Reserved {self.compaction_notice_token_reserve} tokens for " + f"potential notice (compaction budget: {compaction_budget:,}; " + f"governing budget: {budget:,})" + ) # Determine working messages based on whether factory is set if self._system_prompt_factory: @@ -306,11 +298,14 @@ async def get_messages_for_request( token_count = self._estimate_tokens(working_messages) - # Check if compaction needed (using effective budget with notice reserve deducted) - if self._should_compact(token_count, effective_budget): + # Every threshold check uses the canonical integral boundary helper. + if self._should_compact(token_count, compaction_budget): # Compact EPHEMERALLY - returns new list, working_messages unchanged compacted = await self._compact_ephemeral( - effective_budget, working_messages + compaction_budget, working_messages + ) + await self._emit( + _EVT_COMPACTION, dict(self._last_compaction_stats or {}) ) logger.info( f"Ephemeral compaction: {len(working_messages)} -> {len(compacted)} messages for this request" @@ -332,21 +327,30 @@ async def get_messages_for_request( and compacted[notice_index].get("role") == "system" ): notice_index += 1 - compacted.insert( - notice_index, - { - "role": "system", - "content": notice, - "metadata": { - "source": "context-compaction", - "ephemeral": True, - }, + notice_message = { + "role": "system", + "content": notice, + "metadata": { + "source": "context-compaction", + "ephemeral": True, }, - ) - logger.debug( - f"Inserted compaction notice at position {notice_index} (level {level}, " - f"verbosity: {self.compaction_notice_verbosity})" - ) + } + with_notice = list(compacted) + with_notice.insert(notice_index, notice_message) + with_notice_tokens = self._estimate_tokens(with_notice) + if with_notice_tokens <= budget: + compacted = with_notice + logger.debug( + f"Inserted compaction notice at position {notice_index} " + f"(level {level}, verbosity: " + f"{self.compaction_notice_verbosity})" + ) + else: + logger.info( + "Omitting compaction notice because it would exceed " + f"the governing budget ({with_notice_tokens:,} > " + f"{budget:,} tokens)" + ) return compacted @@ -391,7 +395,7 @@ async def should_compact( True if stored context usage >= compact_threshold. """ report = self.usage_report(provider=provider, token_budget=token_budget) - return report["tokens"] >= report["threshold_tokens"] + return self._should_compact(report["tokens"], report["budget"]) def usage_report( self, @@ -411,7 +415,7 @@ def usage_report( "budget": budget, "pct": pct, "threshold": self.compact_threshold, - "threshold_tokens": int(budget * self.compact_threshold), + "threshold_tokens": self._threshold_tokens(budget), "target_tokens": int(budget * self.target_usage), "messages": len(self.messages), } @@ -489,6 +493,7 @@ async def compact( if after_tokens >= before_tokens and after_messages >= before_messages: stats["compacted"] = False stats["reason"] = "no_reduction" + await self._emit(_EVT_COMPACTION, stats) logger.info( "Persistent compaction made no reduction; stored history left unchanged" ) @@ -496,6 +501,8 @@ async def compact( self.messages = compacted stats["compacted"] = True + stats.pop("reason", None) + await self._emit(_EVT_COMPACTION, stats) await self._emit(_EVT_POST_COMPACT, stats) logger.info( f"Persistent compaction committed: {before_messages} -> {len(compacted)} messages, " @@ -513,9 +520,9 @@ async def _emit(self, event: str, payload: dict[str, Any]) -> None: logger.warning(f"Could not emit {event} event: {e}") def _should_compact(self, token_count: int, budget: int) -> bool: - """Check if context should be compacted.""" + """Check compaction using the canonical integral threshold boundary.""" usage = token_count / budget if budget > 0 else 0 - should = usage >= self.compact_threshold + should = budget > 0 and token_count >= self._threshold_tokens(budget) if should: logger.info( f"Context at {usage:.1%} capacity ({token_count:,}/{budget:,} tokens), " @@ -523,6 +530,10 @@ def _should_compact(self, token_count: int, budget: int) -> bool: ) return should + def _threshold_tokens(self, budget: int) -> int: + """Return the first integral token count at or above the threshold.""" + return ceil(budget * self.compact_threshold) if budget > 0 else 0 + async def _compact_ephemeral( self, budget: int, source_messages: list[dict[str, Any]] | None = None ) -> list[dict[str, Any]]: @@ -591,15 +602,16 @@ async def _compact_ephemeral( ] total_tools = len(tool_result_indices) - # Always protect the last N tool results from truncation - protected_tool_indices = set( - tool_result_indices[-self.protected_tool_results :] + # Protect the configured number of results in the recent-history zone. + protected_tool_indices = self._protected_tool_indices( + working_messages, tool_result_indices ) - # Calculate wave boundaries (25% chunks) - wave1_end = int(total_tools * 0.25) - wave2_end = int(total_tools * 0.50) - wave3_end = int(total_tools * 0.75) + # Ceiling keeps small histories progressive too: a single old tool + # result belongs to the first wave rather than bypassing truncation. + wave1_end = ceil(total_tools * 0.25) + wave2_end = ceil(total_tools * 0.50) + wave3_end = ceil(total_tools * 0.75) total_truncated = 0 total_removed = 0 @@ -690,8 +702,8 @@ async def _compact_ephemeral( tool_result_indices = [ i for i, msg in enumerate(working_messages) if msg.get("role") == "tool" ] - protected_tool_indices = set( - tool_result_indices[-self.protected_tool_results :] + protected_tool_indices = self._protected_tool_indices( + working_messages, tool_result_indices ) wave3_start = int(len(tool_result_indices) * 0.50) wave3_end = int(len(tool_result_indices) * 0.75) @@ -753,8 +765,8 @@ async def _compact_ephemeral( tool_result_indices = [ i for i, msg in enumerate(working_messages) if msg.get("role") == "tool" ] - protected_tool_indices = set( - tool_result_indices[-self.protected_tool_results :] + protected_tool_indices = self._protected_tool_indices( + working_messages, tool_result_indices ) truncated, current_tokens = self._truncate_tool_wave( @@ -787,7 +799,10 @@ async def _compact_ephemeral( level7_protection = self.protected_recent * 0.3 working_messages, removed, stubbed, current_tokens = ( self._remove_messages_with_protection( - working_messages, target_tokens, protected_recent=level7_protection + working_messages, + target_tokens, + protected_recent=level7_protection, + allow_user_removal=True, ) ) total_removed += removed @@ -930,16 +945,36 @@ def _truncate_tool_wave( current_tokens = true_total return truncated, current_tokens + def _protected_tool_indices( + self, + messages: list[dict[str, Any]], + tool_result_indices: list[int], + ) -> set[int]: + """Protect recent tool results while leaving old results truncatable. + + Count-based protection is constrained to the configured recent-history + zone. Otherwise a history containing only one very old tool result + would protect it forever and skip the least-destructive truncation phase. + """ + if self.protected_tool_results <= 0: + return set() + recent_boundary = int(len(messages) * (1 - self.protected_recent)) + recent_tools = [i for i in tool_result_indices if i >= recent_boundary] + return set(recent_tools[-self.protected_tool_results :]) + def _remove_messages_with_protection( self, messages: list[dict[str, Any]], target_tokens: int, protected_recent: float, + allow_user_removal: bool = False, ) -> tuple[list[dict[str, Any]], int, int, int]: """ Remove oldest messages with specified protection level. - User messages are NEVER removed - they may be stubbed if still over target. + User messages are preserved at normal levels. At the last-resort level, + intermediate user messages may be removed, while the first and last user + messages remain protected. Returns (new_messages, removed_count, stubbed_count, new_token_count). """ @@ -972,17 +1007,29 @@ def _remove_messages_with_protection( # Always protect the LAST user message (current context) if last_user_idx is not None: protected_indices.add(last_user_idx) + if allow_user_removal and first_user_idx is not None: + protected_indices.add(first_user_idx) # Protect last N% of messages (using the passed protection level) protected_boundary = int(len(messages) * (1 - protected_recent)) for i in range(protected_boundary, len(messages)): protected_indices.add(i) - # Removal candidates exclude ALL user messages (they can only be stubbed, not removed) + # Retain tool results that took the preferred truncation path. The pair + # helpers will then reject removing their corresponding tool calls. + protected_indices.update( + i + for i, msg in enumerate(messages) + if msg.get("role") == "tool" and msg.get("_truncated") + ) + + # Intermediate user messages become candidates only in the final + # last-resort pass. First and last user messages remain protected. removal_candidates = [ i for i in range(len(messages)) - if i not in protected_indices and i not in user_message_indices + if i not in protected_indices + and (i not in user_message_indices or allow_user_removal) ] # Precompute per-message token counts and a tool_call_id -> indices map @@ -1233,9 +1280,6 @@ async def _finalize_compaction_with_stats( } self._last_compaction_stats = stats - # Emit event if hooks available - await self._emit(_EVT_COMPACTION, stats) - return final_messages def _truncate_tool_result(self, msg: dict[str, Any]) -> dict[str, Any]: From dc7a7a1d34a74ab03264634c3b952b19d7b2d2bf Mon Sep 17 00:00:00 2001 From: Sam Schillace Date: Thu, 6 Aug 2026 17:24:59 -0700 Subject: [PATCH 4/4] fix: enforce strict token-monotonic compaction --- amplifier_module_context_simple/__init__.py | 11 ++-- tests/test_persistent_compaction.py | 59 +++++++++++++++++++++ 2 files changed, 66 insertions(+), 4 deletions(-) diff --git a/amplifier_module_context_simple/__init__.py b/amplifier_module_context_simple/__init__.py index 1ec4be4..43709a3 100644 --- a/amplifier_module_context_simple/__init__.py +++ b/amplifier_module_context_simple/__init__.py @@ -490,7 +490,7 @@ async def compact( stats["persistent"] = True stats["forced"] = force - if after_tokens >= before_tokens and after_messages >= before_messages: + if after_tokens >= before_tokens: stats["compacted"] = False stats["reason"] = "no_reduction" await self._emit(_EVT_COMPACTION, stats) @@ -932,14 +932,17 @@ def _truncate_tool_wave( if msg.get("role") != "tool": # Verify it's still a tool message continue if not msg.get("_truncated"): + replacement = self._truncate_tool_result(msg) + old_len = len(str(msg)) // 4 + new_len = len(str(replacement)) // 4 + if new_len >= old_len: + continue if true_total is None: # First mutation in this call: establish the true baseline # once (matches what _estimate_tokens(messages) would have # returned immediately before this mutation). true_total = self._estimate_tokens(messages) - old_len = len(str(msg)) // 4 - messages[i] = self._truncate_tool_result(msg) - new_len = len(str(messages[i])) // 4 + messages[i] = replacement true_total += new_len - old_len truncated += 1 current_tokens = true_total diff --git a/tests/test_persistent_compaction.py b/tests/test_persistent_compaction.py index a2c8df2..56c369f 100644 --- a/tests/test_persistent_compaction.py +++ b/tests/test_persistent_compaction.py @@ -129,3 +129,62 @@ async def test_usage_report_shape(): assert report["tokens"] > 0 assert 0.0 <= report["pct"] assert report["messages"] == len(context.messages) + + +def test_tool_truncation_rejects_token_growth(): + context = SimpleContextManager(truncate_chars=374) + message = { + "role": "tool", + "content": "x" * 411, + "tool_call_id": "call-1", + } + messages = [message] + before_tokens = context._estimate_tokens(messages) + candidate = context._truncate_tool_result(message) + + assert before_tokens == 117 + assert context._estimate_tokens([candidate]) == 133 + + truncated, after_tokens = context._truncate_tool_wave( + messages, + indices=[0], + protected_indices=set(), + target_tokens=0, + current_tokens=before_tokens, + ) + + assert truncated == 0 + assert after_tokens == before_tokens + assert messages == [message] + + +@pytest.mark.asyncio +async def test_compact_rejects_fewer_messages_when_tokens_grow(monkeypatch): + hooks = _RecordingHooks() + context = SimpleContextManager( + max_tokens=10, + compact_threshold=0.9, + target_usage=0.5, + hooks=hooks, + ) + context.messages = [ + {"role": "user", "content": "first"}, + {"role": "assistant", "content": "second"}, + ] + original = list(context.messages) + candidate = [{"role": "user", "content": "x" * 200}] + assert len(candidate) < len(original) + assert context._estimate_tokens(candidate) > context._estimate_tokens(original) + + async def growing_candidate(_budget, _messages): + return candidate + + monkeypatch.setattr(context, "_compact_ephemeral", growing_candidate) + + stats = await context.compact(force=True) + + assert stats["compacted"] is False + assert stats["reason"] == "no_reduction" + assert stats["after_tokens"] > stats["before_tokens"] + assert context.messages == original + assert "context:post_compact" not in [event for event, _ in hooks.events]