From 4ea79443bc8e7349adb1f6516a3da3fae8225cd6 Mon Sep 17 00:00:00 2001
From: Teknium <127238744+teknium1@users.noreply.github.com>
Date: Sun, 1 Nov 2026 22:17:50 -0800
Subject: [PATCH] perf(agent): the token memo keeps no more idle entries than
 it always had

#147819 raised the limits to 65,536 entries / 64 MiB, so closed sessions could hold more than main
said 4,096 entries ever did. Both limits only evict entries idle for 60 s, so lowering them to 4,096 /
16 MiB costs a running conversation nothing.
---
 agent/model_metadata_token_memo.py        | 13 +++++++------
 tests/agent/test_token_memo_lru_budget.py | 11 +++++++++++
 2 files changed, 18 insertions(+), 6 deletions(-)

diff --git a/agent/model_metadata_token_memo.py b/agent/model_metadata_token_memo.py
index 36012fc94e..79c2334313 100644
--- a/agent/model_metadata_token_memo.py
+++ b/agent/model_metadata_token_memo.py
@@ -3,10 +3,11 @@
 Each entry pins every string its fingerprint holds (the pin keeps the ``id()`` key sound), so the memo is
 bounded two ways, both least-recently-used first so a conversation that is still running never pays:
 
-* a retained-bytes budget: pinned string bytes plus the fingerprint structure around them. A real-shaped
-  1M-token history retains about 9 MB, so the default keeps several maxed conversations hot at once while
-  the dead sessions of a long-lived process age out;
-* an entry count, as a backstop for the bytes estimate.
+* a retained-bytes budget: pinned string bytes plus the fingerprint structure around them;
+* an entry count, the same 4,096 the memo always had.
+
+Both bounds only govern IDLE entries (below), so they are set to hold no more than main's 4,096-entry memo
+ever did: what a running conversation needs is protected by the active window, never by the budget.
 
 Neither bound ever evicts an entry looked up in the last ``_ACTIVE_WINDOW_S`` seconds: a conversation larger
 than the budget (image-heavy history) overshoots it while it runs instead of missing on every message, and the
@@ -28,8 +29,8 @@ from typing import Any
 # fingerprint -> (pins, text tokens, image count). Order is recency: a hit moves its entry to the end, and an
 # OrderedDict evicts from the front in O(1) where a plain dict rescans its deleted slots on every eviction.
 _MSG_TOKENS_CACHE: OrderedDict[Any, tuple[list, int, int]] = OrderedDict()
-_MSG_TOKENS_CACHE_MAX = 65536
-_MSG_TOKENS_CACHE_MAX_BYTES = 64 * 1024 * 1024
+_MSG_TOKENS_CACHE_MAX = 4096
+_MSG_TOKENS_CACHE_MAX_BYTES = 16 * 1024 * 1024
 # Measured on real histories: the fingerprint tuple, its id int and the pin slot cost ~120 B per string leaf.
 _LEAF_OVERHEAD_BYTES = 120
 # An entry looked up this recently is never evicted (see module docstring).
diff --git a/tests/agent/test_token_memo_lru_budget.py b/tests/agent/test_token_memo_lru_budget.py
index b0792ced5a..b7ec14419f 100644
--- a/tests/agent/test_token_memo_lru_budget.py
+++ b/tests/agent/test_token_memo_lru_budget.py
@@ -116,3 +116,14 @@ def test_a_hit_racing_eviction_never_makes_it_drop_a_different_active_entry(monk
     memo.remember_message_tokens("new", (memo.MemoPins(["new text"]), 1, 0))
     assert reader_done.wait(5)
     assert "active" in cache and "new" in cache
+
+
+def test_idle_entries_never_outnumber_the_memo_main_always_had(monkeypatch):
+    """Closed sessions' entries (idle) stay within the original 4,096-entry footprint; only entries a running
+    conversation touched in the last minute may exceed it."""
+    mm._MSG_TOKENS_CACHE.clear()
+    now = _fake_clock(monkeypatch)
+    for i in range(6000):
+        now[0] += 61  # each message's session is idle before the next one arrives
+        mm.estimate_messages_tokens_rough([{"role": "tool", "content": f"closed-{i} " + "q" * 300}])
+    assert len(mm._MSG_TOKENS_CACHE) <= 4096