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:
parent
cf803603dc
commit
e8656bedfe
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
Loading…
Reference in New Issue