From e8656bedfede4e519b54159c435dcecdbf571351 Mon Sep 17 00:00:00 2001
From: kshitij <82637225+kshitijk4poor@users.noreply.github.com>
Date: Sun, 2 Aug 2026 20:02:32 +0530
Subject: [PATCH] fix(agent): byte budget for the canon-args memo
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Review follow-up on the #76098 salvage: the 4096-entry count bound alone
doesn't bound MEMORY — write_file/patch argument strings run 100KB+, so
a long-lived gateway process under sustained heavy write workloads could
pin ~800MB of evicted-session strings. A 32MB byte budget extends the
existing FIFO eviction; common-case args (0.5-2KB) never hit it. New
guard test mutation-checked (fails with the byte leg disabled).
---
 agent/conversation_loop.py                 | 17 ++++++++++++--
 tests/agent/test_canon_args_memo_parity.py | 27 ++++++++++++++++++----
 2 files changed, 38 insertions(+), 6 deletions(-)

diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py
index 8c602a16c6ae5..68453df309f77 100644
--- a/agent/conversation_loop.py
+++ b/agent/conversation_loop.py
@@ -754,6 +754,12 @@ _CONTENT_POLICY_RECOVERY_HINT = (
 # _MSG_TOKENS_CACHE idiom in agent/model_metadata.py.
 _CANON_ARGS_CACHE: Dict[str, str] = {}
 _CANON_ARGS_CACHE_MAX = 4096
+# Count bound alone doesn't bound MEMORY: write_file/patch argument strings
+# run 100KB+, so 4096 entries could pin ~800MB in a long-lived gateway
+# process. The byte budget keeps the memo effective for the common case
+# (args ~0.5-2KB) while bounding the worst case.
+_CANON_ARGS_CACHE_MAX_BYTES = 32 * 1024 * 1024
+_canon_args_cache_bytes = 0
 
 
 def _canonicalize_tool_call_arguments(arg_str: str) -> str:
@@ -762,6 +768,7 @@ def _canonicalize_tool_call_arguments(arg_str: str) -> str:
     Raises whatever ``json.loads`` raises on malformed input; the caller
     falls back to ``_repair_tool_call_arguments``, exactly as before.
     """
+    global _canon_args_cache_bytes
     cached = _CANON_ARGS_CACHE.get(arg_str)
     if cached is not None:
         return cached
@@ -769,9 +776,15 @@ def _canonicalize_tool_call_arguments(arg_str: str) -> str:
         json.loads(arg_str), separators=(",", ":"), sort_keys=True,
     )
     _CANON_ARGS_CACHE[arg_str] = canonical
-    while len(_CANON_ARGS_CACHE) > _CANON_ARGS_CACHE_MAX:
+    _canon_args_cache_bytes += len(arg_str) + len(canonical)
+    while len(_CANON_ARGS_CACHE) > _CANON_ARGS_CACHE_MAX or (
+        _canon_args_cache_bytes > _CANON_ARGS_CACHE_MAX_BYTES
+        and len(_CANON_ARGS_CACHE) > 1
+    ):
         try:
-            _CANON_ARGS_CACHE.pop(next(iter(_CANON_ARGS_CACHE)))
+            evicted_key = next(iter(_CANON_ARGS_CACHE))
+            evicted_val = _CANON_ARGS_CACHE.pop(evicted_key)
+            _canon_args_cache_bytes -= len(evicted_key) + len(evicted_val)
         except (StopIteration, KeyError, RuntimeError):
             break
     return canonical
diff --git a/tests/agent/test_canon_args_memo_parity.py b/tests/agent/test_canon_args_memo_parity.py
index d35e01f2c52cd..827c58eeae1ae 100644
--- a/tests/agent/test_canon_args_memo_parity.py
+++ b/tests/agent/test_canon_args_memo_parity.py
@@ -39,11 +39,30 @@ def _clear_canon_cache():
     # setup on the pre-fix tree, so sabotage runs record real test FAILURES
     # (AttributeError inside each test body) instead of collection errors.
     cache = getattr(cl, "_CANON_ARGS_CACHE", None)
-    if cache is not None:
-        cache.clear()
+
+    def _reset():
+        if cache is not None:
+            cache.clear()
+        if hasattr(cl, "_canon_args_cache_bytes"):
+            cl._canon_args_cache_bytes = 0
+
+    _reset()
     yield
-    if cache is not None:
-        cache.clear()
+    _reset()
+
+
+def test_cache_bounded_by_bytes():
+    """Large argument strings (write_file contents run 100KB+) must not pin
+    unbounded memory: the byte budget evicts before the count bound."""
+    big = json.dumps({"path": "/tmp/big.py", "content": "y" * 200_000})
+    for i in range(300):  # 300 x ~400KB (key+value) >> 32MB budget
+        cl._canonicalize_tool_call_arguments(
+            big[:-1] + f',"n":{i}}}'
+        )
+    assert cl._canon_args_cache_bytes <= cl._CANON_ARGS_CACHE_MAX_BYTES, (
+        f"cache holds {cl._canon_args_cache_bytes} bytes — byte budget "
+        "regressed; large tool-call args pin unbounded memory again")
+    assert len(cl._CANON_ARGS_CACHE) >= 1  # still memoizes something
 
 
 def build_history(n_tool_calls, arg_bytes=2048):