fix(agent): compression releases the message bodies its summary replaced from the token memo

close() forgets only the memo entries pinned by the history it still holds, so bodies a compression
already summarized away stayed pinned for the process's life (Greptile on #137087). The compression
facade now forgets the pre-compression list minus whatever the compressed list still carries.
This commit is contained in:
kshitijk4poor authored and kshitij committed 2026-10-20 18:08:50 +05:30
1 parent 1ee55ac331
commit a83bfc453f
3 files changed
+55 -9

No files matched your search

+6
View File
@@ -198,6 +198,12 @@ def _mirror_result_onto_live_lists(agent, result, messages, *, direct_path: bool
if isinstance(session_messages, list) and session_messages is not messages:
# Durable-parent adoption can leave `_session_messages` on the pre-adoption list.
_sync_persisted_markers(session_messages, result_messages)
if result_messages is not messages:
# The token-estimate memo pins message strings; bodies the summary replaced can never hit again,
# and close() only sees the compressed history.
from agent.model_metadata_token_memo import forget_message_token_memo
with contextlib.suppress(Exception):
forget_message_token_memo(messages, keep=result_messages)
def _rebind_caller_session_context(agent) -> None:
+17 -9
View File
@@ -10,14 +10,8 @@ from __future__ import annotations
from typing import Any
def forget_message_token_memo(messages: Any) -> int:
"""Evict memo entries pinning any long string value of ``messages``; return how many.
Eviction never changes an estimate: a later lookup recomputes the same value.
"""
from agent.model_metadata import _MSG_TOKENS_CACHE
owned: set[int] = set()
def _long_string_ids(messages: Any) -> set[int]:
ids: set[int] = set()
stack = [messages]
while stack:
value = stack.pop()
@@ -26,7 +20,21 @@ def forget_message_token_memo(messages: Any) -> int:
elif type(value) is list or type(value) is tuple:
stack.extend(value)
elif type(value) is str and len(value) > 64: # short values ("assistant", tool names) are shared
owned.add(id(value))
ids.add(id(value))
return ids
def forget_message_token_memo(messages: Any, *, keep: Any = None) -> int:
"""Evict memo entries pinning any long string value of ``messages`` that ``keep`` does not also hold;
return how many.
Eviction never changes an estimate: a later lookup recomputes the same value.
"""
from agent.model_metadata import _MSG_TOKENS_CACHE
owned = _long_string_ids(messages)
if owned and keep is not None:
owned -= _long_string_ids(keep)
if not owned:
return 0
stale = [key for key, (pins, _text, _images) in list(_MSG_TOKENS_CACHE.items())
@@ -42,3 +42,35 @@ def test_close_evicts_the_conversations_memo_entries_and_keeps_others():
assert id(other) in _pinned_ids()
# Eviction only drops a pure-function memo: the estimate is unchanged when recomputed.
assert estimate_messages_tokens_rough(closed) == before
def test_compression_evicts_the_bodies_the_summary_replaced(tmp_path, monkeypatch):
"""close() only sees the compressed history, so the bodies compression dropped are released then."""
import os
from unittest.mock import MagicMock, patch
from agent import conversation_compression
from agent.conversation_compression import CompressionCommitFence
from hermes_state import SessionDB
db = SessionDB(db_path=tmp_path / "state.db")
db.create_session("S", source="cli")
with patch.dict(os.environ, {"OPENROUTER_API_KEY": "test-key"}):
agent = AIAgent(api_key="test-key", base_url="https://openrouter.ai/api/v1", model="test/model",
platform="cli", quiet_mode=True, session_db=db, session_id="S",
skip_context_files=True, skip_memory=True)
agent.context_compressor = MagicMock()
dropped = "pre-compression tool output " * 5000
kept = "kept tail reply " * 5000
history = [{"role": "user", "content": "q" * 80}, {"role": "tool", "content": dropped},
{"role": "assistant", "content": kept}]
estimate_messages_tokens_rough(history)
compressed = [{"role": "user", "content": "[summary] " + "s" * 80}, history[2]]
monkeypatch.setattr(conversation_compression, "compress_context",
lambda *_a, **_kw: (compressed, "prompt"))
assert id(dropped) in _pinned_ids()
agent._compress_context(history, "sys", commit_fence=CompressionCommitFence())
assert id(dropped) not in _pinned_ids()
assert id(kept) in _pinned_ids() # still live in the compressed history