fix(agent): byte budget for the canon-args memo

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).
This commit is contained in:
kshitij 2026-08-02 20:02:32 +05:30
parent cf803603dc
commit e8656bedfe
2 changed files with 38 additions and 6 deletions

View File

@ -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

View File

@ -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):