diff --git a/agent/agent_runtime_helpers.py b/agent/agent_runtime_helpers.py index 6192556d2cb..5505d707562 100644 --- a/agent/agent_runtime_helpers.py +++ b/agent/agent_runtime_helpers.py @@ -249,12 +249,42 @@ def sanitize_tool_call_arguments( *, logger=None, session_id: str = None, + cursor: Optional[dict] = None, ) -> int: - """Repair corrupted assistant tool-call argument JSON in-place.""" + """Repair corrupted assistant tool-call argument JSON in-place. + + ``cursor`` (optional) is a caller-owned dict used to skip re-validating + messages already validated on a previous call. It stores, under + ``"prefix"``, the exact message *objects* (strong references) validated + last time, in order. On the next call, the longest contiguous prefix of + ``messages`` whose objects are ``is``-identical to the stored prefix is + skipped; scanning starts at the first divergence (conservative: any + reordering, truncation, compression rewrite, or mid-list insertion breaks + identity at that index and everything from there is re-scanned). + + Safety argument for skipping: a message in the matched prefix was fully + scanned before — every tool_call argument was either already valid JSON + or was rewritten to ``"{}"`` (valid). The only code paths that mutate + ``function["arguments"]`` on live history dicts between calls are the + surrogate / non-ASCII sanitizers, which substitute characters *inside* + JSON string values and cannot invalidate JSON syntax. Compression, + repair, undo, and steer paths replace or reorder message dicts, which + breaks the identity match and forces a re-scan. Holding strong + references (the objects themselves, not ``id()``s) makes address reuse + aliasing (#50372-style) impossible. + """ log = logger or logging.getLogger(__name__) if not isinstance(messages, list): return 0 + start_index = 0 + if cursor is not None: + prev_prefix = cursor.get("prefix") + if isinstance(prev_prefix, list): + limit = min(len(prev_prefix), len(messages)) + while start_index < limit and messages[start_index] is prev_prefix[start_index]: + start_index += 1 + repaired = 0 marker = _ra().AIAgent._TOOL_CALL_ARGUMENTS_CORRUPTION_MARKER @@ -275,7 +305,7 @@ def sanitize_tool_call_arguments( existing_text = str(existing) tool_msg["content"] = f"{marker}\n{existing_text}" - message_index = 0 + message_index = start_index while message_index < len(messages): msg = messages[message_index] if not isinstance(msg, dict) or msg.get("role") != "assistant": @@ -356,6 +386,12 @@ def sanitize_tool_call_arguments( message_index += 1 + if cursor is not None: + # Strong references to the exact objects validated this call, in + # order. Any future divergence (compression, undo, repair, steer) + # breaks identity at the divergent index and re-scans from there. + cursor["prefix"] = messages[:] + return repaired diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 95c4f7ea925..5d96b8e6512 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -1383,10 +1383,23 @@ def run_conversation( # However, providers like Moonshot AI require a separate 'reasoning_content' field # on assistant messages with tool_calls. We handle both cases here. request_logger = getattr(agent, "logger", None) or logging.getLogger(__name__) + # Per-agent validation cursor: skips re-json.loads-ing tool_call + # arguments on history messages already validated in a previous + # iteration. Identity-keyed (strong refs) — compression/undo/repair + # rewriting the list breaks the prefix match and forces a re-scan + # from the divergence point. See sanitize_tool_call_arguments. + _sanitize_cursor = getattr(agent, "_sanitize_args_cursor", None) + if _sanitize_cursor is None: + _sanitize_cursor = {} + try: + agent._sanitize_args_cursor = _sanitize_cursor + except Exception: + pass repaired_tool_calls = agent._sanitize_tool_call_arguments( messages, logger=request_logger, session_id=agent.session_id, + cursor=_sanitize_cursor, ) if repaired_tool_calls > 0: request_logger.info( diff --git a/agent/model_metadata.py b/agent/model_metadata.py index 6aacc71b154..99f18d8851a 100644 --- a/agent/model_metadata.py +++ b/agent/model_metadata.py @@ -2851,14 +2851,92 @@ def estimate_messages_tokens_rough(messages: List[Dict[str, Any]]) -> int: image — the Anthropic pricing model — instead of counting raw base64 character length. Without this, a single ~1MB screenshot would be estimated at ~250K tokens and trigger premature context compression. + + Per-message results are memoized (see ``_estimate_message_tokens_cached``) + keyed on a deep *identity fingerprint* of the message, so re-walking a + long history every iteration only pays for messages whose object graph + actually changed. The memo is exact: equal fingerprints imply identical + leaf objects and structure, hence an identical estimate. """ _IMAGE_TOKEN_COST = 1500 - text_tokens = 0 - image_tokens = 0 + total = 0 for msg in messages: - text_tokens += _estimate_message_tokens_without_images(msg) - image_tokens += _count_image_tokens(msg, _IMAGE_TOKEN_COST) - return text_tokens + image_tokens + total += _estimate_message_tokens_cached(msg, _IMAGE_TOKEN_COST) + return total + + +# --- Per-message token-estimate memo ------------------------------------- +# +# ``estimate_messages_tokens_rough`` is called on the full history every +# loop iteration (conversation_loop preflight), repeatedly during compaction +# telemetry, and inside an O(n^2) shrink loop in moa_loop. The per-message +# helpers are pure functions of the message's value, so a memo keyed on a +# fingerprint that uniquely determines the value is exactly equivalent. +# +# Fingerprint design (soundness argument): +# * strings are fingerprinted by ``id()`` AND pinned (a strong reference is +# stored in the cache entry). While the entry lives, that id cannot be +# reused by another object, so id-equality implies object-equality — +# strings are immutable, so value-equality too (no #50372-style aliasing). +# * ints/floats/bools/None are fingerprinted by value. +# * dicts/lists recurse structurally, preserving key order — ``str(shadow)`` +# depends on insertion order, so order is part of the key. +# * any other type aborts the memo and falls through to a direct compute. +# Equal fingerprints therefore imply deep-equal messages built from identical +# immutable leaves ⇒ identical ``str(shadow)`` bytes ⇒ identical estimate. +# +# Because the api_messages build shallow-copies history dicts each iteration, +# the copies share the same content strings — so unchanged history messages +# hit the memo even though the outer dicts are fresh objects every turn. +_MSG_TOKENS_CACHE: Dict[Any, Tuple[list, int]] = {} +_MSG_TOKENS_CACHE_MAX = 4096 + + +def _msg_fingerprint(value: Any, pins: list) -> Any: + if value is None or value is True or value is False: + return value + t = type(value) + if t is str: + pins.append(value) + return ("s", id(value)) + if t is int or t is float: + return ("n", t.__name__, value) + if t is dict: + return ("d", tuple( + (_msg_fingerprint(k, pins), _msg_fingerprint(v, pins)) + for k, v in value.items() + )) + if t is list: + return ("l", tuple(_msg_fingerprint(v, pins) for v in value)) + if t is tuple: + return ("t", tuple(_msg_fingerprint(v, pins) for v in value)) + raise ValueError("unfingerprintable message value") + + +def _estimate_message_tokens_cached(msg: Any, image_cost: int) -> int: + try: + pins: list = [] + key = _msg_fingerprint(msg, pins) + hash(key) + except Exception: + return ( + _estimate_message_tokens_without_images(msg) + + _count_image_tokens(msg, image_cost) + ) + cached = _MSG_TOKENS_CACHE.get(key) + if cached is not None: + return cached[1] + tokens = ( + _estimate_message_tokens_without_images(msg) + + _count_image_tokens(msg, image_cost) + ) + _MSG_TOKENS_CACHE[key] = (pins, tokens) + while len(_MSG_TOKENS_CACHE) > _MSG_TOKENS_CACHE_MAX: + try: + _MSG_TOKENS_CACHE.pop(next(iter(_MSG_TOKENS_CACHE))) + except (StopIteration, KeyError, RuntimeError): + break + return tokens def _count_image_tokens(msg: Dict[str, Any], cost_per_image: int) -> int: diff --git a/run_agent.py b/run_agent.py index 641abaf6f6c..dde8a8797cc 100644 --- a/run_agent.py +++ b/run_agent.py @@ -2004,7 +2004,27 @@ class AIAgent: if isinstance(item, dict) } - for _msg_idx, msg in enumerate(messages): + # Bounded scan: skip the longest identity-matched prefix of the + # list snapshot taken at the end of the previous successful flush. + # Every message in that snapshot was already given its final + # disposition (written+stamped, stamped as durable history, or + # skipped as ephemeral scaffolding / non-dict), and no code path + # pops _DB_PERSISTED_MARKER from a live dict in place (compression + # strips markers on fresh copies, which breaks identity here and + # forces a full re-scan). Identity match ⇒ identical skip decision, + # so starting after the matched prefix is behavior-preserving. + _scan_start = 0 + _prev_prefix = getattr(self, "_db_flush_scan_prefix", None) + if isinstance(_prev_prefix, list): + _limit = min(len(_prev_prefix), len(messages)) + while ( + _scan_start < _limit + and messages[_scan_start] is _prev_prefix[_scan_start] + ): + _scan_start += 1 + + for _msg_idx in range(_scan_start, len(messages)): + msg = messages[_msg_idx] if not isinstance(msg, dict): continue # Never write ephemeral recovery scaffolding to the session @@ -2153,8 +2173,14 @@ class AIAgent: # allocated next turn at a recycled address. self._flushed_db_message_ids = set() self._last_flushed_db_idx = len(messages) + # Snapshot for the bounded scan above — only on full success, so + # a partially-processed list can never be treated as settled. + self._db_flush_scan_prefix = messages[:] return True except Exception as e: + # Force a full re-scan on the next flush: an exception mid-loop + # leaves messages with mixed dispositions. + self._db_flush_scan_prefix = None logger.warning("Session DB append_message failed: %s", e) return False @@ -6738,10 +6764,13 @@ class AIAgent: *, logger=None, session_id: str = None, + cursor=None, ) -> int: """Forwarder — see ``agent.agent_runtime_helpers.sanitize_tool_call_arguments``.""" from agent.agent_runtime_helpers import sanitize_tool_call_arguments - return sanitize_tool_call_arguments(messages, logger=logger, session_id=session_id) + return sanitize_tool_call_arguments( + messages, logger=logger, session_id=session_id, cursor=cursor + ) def _should_sanitize_tool_calls(self) -> bool: """Determine if tool_calls need sanitization for strict APIs. diff --git a/tests/agent/test_cursor_optimizations_parity.py b/tests/agent/test_cursor_optimizations_parity.py new file mode 100644 index 00000000000..c08b012aca3 --- /dev/null +++ b/tests/agent/test_cursor_optimizations_parity.py @@ -0,0 +1,269 @@ +"""Byte-parity + benchmark harness for the per-iteration cursor optimizations. + +Drives the pure functions directly (no AIAgent): + 1. sanitize_tool_call_arguments (with/without cursor) + 2. estimate_messages_tokens_rough (memoized) vs a reference reimplementation + 3. _flush_messages_to_session_db bounded scan — simulated via a stub agent + +Run: HERMES worktree venv python parity_harness.py +""" +import copy +import json +import random +import statistics +import sys +import time + +sys.path.insert(0, ".") + +random.seed(1234) + +UNI = "日本語テキスト🎉 café Ω ≈ 中文字符串" + + +def build_history(n): + """Synthetic conversation: user/assistant/tool cycles, malformed args, unicode.""" + msgs = [] + i = 0 + while len(msgs) < n: + msgs.append({"role": "user", "content": f"question {i} {UNI} " + "x" * random.randint(10, 400)}) + if i % 3 == 0: + args = json.dumps({"q": f"val {i}", "u": UNI, "n": i}) + if i % 9 == 0: + args = '{"broken": tru' # malformed + elif i % 6 == 0: + args = "" # empty + msgs.append({ + "role": "assistant", "content": "", + "tool_calls": [{"id": f"call_{i}", "type": "function", + "function": {"name": "web_search", "arguments": args}}], + }) + msgs.append({"role": "tool", "tool_call_id": f"call_{i}", + "name": "web_search", "content": f"result {i} {UNI}"}) + else: + msgs.append({"role": "assistant", "content": f"answer {i} " + "y" * random.randint(10, 600), + "reasoning_content": f"thinking {i}"}) + if i % 7 == 0 and msgs: + msgs[-1]["content"] = [{"type": "text", "text": f"part {i}"}, + {"type": "image_url", "image_url": {"url": "data:image/png;base64,AAAA"}}] + i += 1 + return msgs[:n] + + +# ---------- reference (pre-optimization) implementations ---------- + +from agent.model_metadata import ( + estimate_messages_tokens_rough, + _estimate_message_tokens_without_images, + _count_image_tokens, + _MSG_TOKENS_CACHE, +) + + +def estimate_messages_tokens_rough_OLD(messages): + _IMAGE_TOKEN_COST = 1500 + text_tokens = 0 + image_tokens = 0 + for msg in messages: + text_tokens += _estimate_message_tokens_without_images(msg) + image_tokens += _count_image_tokens(msg, _IMAGE_TOKEN_COST) + return text_tokens + image_tokens + + +from agent.agent_runtime_helpers import sanitize_tool_call_arguments + + +def simulate_compression(msgs): + """Rewrite the middle of the history with fresh dict copies + a summary.""" + head, mid, tail = msgs[:2], msgs[2:-6], msgs[-6:] + summary = {"role": "user", "content": "SUMMARY OF DROPPED CONTEXT " + UNI} + new = [dict(m) if isinstance(m, dict) else m for m in head] + new.append(summary) + new.extend(dict(m) if isinstance(m, dict) else m for m in tail) + msgs[:] = new + + +def test_parity_sanitize_cursor(): + print("=== parity: sanitize_tool_call_arguments cursor ===") + for n in (50, 200, 500): + base = build_history(n) + old_list = copy.deepcopy(base) + new_list = copy.deepcopy(base) + cursor = {} + for iteration in range(3): + r_old = sanitize_tool_call_arguments(old_list) + r_new = sanitize_tool_call_arguments(new_list, cursor=cursor) + assert r_old == r_new, (n, iteration, r_old, r_new) + assert old_list == new_list, f"list mismatch n={n} it={iteration}" + # append a new exchange (one malformed) between iterations + for lst in (old_list, new_list): + lst.append({"role": "assistant", "content": "", + "tool_calls": [{"id": f"c{iteration}", "type": "function", + "function": {"name": "t", "arguments": '{"bad": '}}]}) + lst.append({"role": "tool", "tool_call_id": f"c{iteration}", "content": "ok"}) + if iteration == 1: + simulate_compression(old_list) + simulate_compression(new_list) + # after mutations, one more full compare + r_old = sanitize_tool_call_arguments(old_list) + r_new = sanitize_tool_call_arguments(new_list, cursor=cursor) + assert r_old == r_new and old_list == new_list + print(f" n={n}: OK (element-wise equal across 3 iterations + compression)") + + +def test_parity_token_memo(): + print("=== parity: estimate_messages_tokens_rough memo ===") + for n in (50, 200, 500): + msgs = build_history(n) + _MSG_TOKENS_CACHE.clear() + for iteration in range(3): + # simulate api_messages copies each iteration (shallow copies) + api = [m.copy() for m in msgs] + old = estimate_messages_tokens_rough_OLD(api) + new = estimate_messages_tokens_rough(api) + assert old == new, (n, iteration, old, new) + msgs.append({"role": "user", "content": f"followup {iteration} {UNI}"}) + if iteration == 1: + simulate_compression(msgs) + # mutate a string in place-ish: replace content of an existing dict + msgs[0]["content"] = "EDITED " + UNI + api = [m.copy() for m in msgs] + assert estimate_messages_tokens_rough_OLD(api) == estimate_messages_tokens_rough(api) + # odd types fall through the memo + weird = [{"role": "user", "content": {"_multimodal": True, "text_summary": "s"}}, + {"role": "user", "content": None}, "not-a-dict", + {"role": "tool", "content": [{"type": "text", "text": UNI}, "raw"], "meta": (1, 2)}] + assert estimate_messages_tokens_rough_OLD(weird) == estimate_messages_tokens_rough(weird) + print(f" n={n}: OK (equal across 3 iterations + compression + in-place edit + odd types)") + + +def test_parity_persist_bounded_scan(): + print("=== parity: _flush_messages_to_session_db bounded scan ===") + import run_agent as ra + + class FakeDB: + def __init__(self): + self.rows = [] + def append_message(self, **kw): + self.rows.append({k: copy.deepcopy(v) for k, v in kw.items()}) + + def make_agent(bounded): + a = ra.AIAgent.__new__(ra.AIAgent) + a.session_id = "s1" + a._session_db = FakeDB() + a._session_db_created = True + a._last_flushed_db_idx = 0 + a._flushed_db_message_ids = set() + a._persist_disabled = False + a._session_persist_lock = None + if not bounded: + # neutralize the cursor: force full scan every time + a._db_flush_scan_prefix = None + return a + + for n in (50, 200, 500): + base = build_history(n) + la, lb = copy.deepcopy(base), copy.deepcopy(base) + A, B = make_agent(False), make_agent(True) + for iteration in range(3): + A._db_flush_scan_prefix = None # baseline: always full scan + ra_ok = A._flush_messages_to_session_db_unlocked(la, None) + rb_ok = B._flush_messages_to_session_db_unlocked(lb, None) + assert ra_ok is True and rb_ok is True + assert A._session_db.rows == B._session_db.rows, f"rows diverge n={n} it={iteration}" + assert la == lb + for lst in (la, lb): + lst.append({"role": "user", "content": f"turn {iteration} {UNI}"}) + lst.append({"role": "assistant", "content": f"reply {iteration}", + "_empty_recovery_synthetic": iteration == 0}) # scaffolding once + if iteration == 1: + # compression-style rewrite: fresh copies without markers + for lst in (la, lb): + head = [dict(m) for m in lst[:3]] + for m in head: + m.pop(ra._DB_PERSISTED_MARKER, None) + tail = [dict(m) for m in lst[-4:]] + for m in tail: + m.pop(ra._DB_PERSISTED_MARKER, None) + lst[:] = head + [{"role": "user", "content": "SUMMARY"}] + tail + A._db_flush_scan_prefix = None + A._flush_messages_to_session_db_unlocked(la, None) + B._flush_messages_to_session_db_unlocked(lb, None) + assert A._session_db.rows == B._session_db.rows and la == lb + print(f" n={n}: OK (identical DB rows + marker stamps across 3 flushes + compression rewrite)") + + +def bench(): + print("=== benchmarks (median of 5, per call) ===") + + def timeit(fn, reps=5): + ts = [] + for _ in range(reps): + t0 = time.perf_counter() + fn() + ts.append(time.perf_counter() - t0) + return statistics.median(ts) * 1e3 # ms + + for n in (50, 200, 500): + msgs = build_history(n) + sanitize_tool_call_arguments(msgs) # settle repairs first + + # sanitize: old (no cursor) vs new (warm cursor) + old_ms = timeit(lambda: sanitize_tool_call_arguments(msgs)) + cur = {} + sanitize_tool_call_arguments(msgs, cursor=cur) + new_ms = timeit(lambda: sanitize_tool_call_arguments(msgs, cursor=cur)) + + # tokens: old walk vs warm memo (on fresh shallow copies, like api_messages) + api = [m.copy() for m in msgs] + told = timeit(lambda: estimate_messages_tokens_rough_OLD([m.copy() for m in msgs])) + _MSG_TOKENS_CACHE.clear() + estimate_messages_tokens_rough([m.copy() for m in msgs]) # warm + tnew = timeit(lambda: estimate_messages_tokens_rough([m.copy() for m in msgs])) + + # persist scan: fully-flushed list, old full walk vs bounded skip + import run_agent as ra + flushed = copy.deepcopy(msgs) + for m in flushed: + if isinstance(m, dict): + m[ra._DB_PERSISTED_MARKER] = True + + def old_scan(): + for _idx, m in enumerate(flushed): + if not isinstance(m, dict): + continue + if ra._is_ephemeral_scaffolding(m): + continue + if m.get(ra._DB_PERSISTED_MARKER): + continue + + prefix = flushed[:] + + def new_scan(): + s = 0 + lim = min(len(prefix), len(flushed)) + while s < lim and flushed[s] is prefix[s]: + s += 1 + for _idx in range(s, len(flushed)): + m = flushed[_idx] + if not isinstance(m, dict): + continue + if ra._is_ephemeral_scaffolding(m): + continue + if m.get(ra._DB_PERSISTED_MARKER): + continue + + pold = timeit(old_scan) + pnew = timeit(new_scan) + + print(f" n={n:3d}: sanitize {old_ms:.3f}ms -> {new_ms:.3f}ms | " + f"tokens {told:.3f}ms -> {tnew:.3f}ms | " + f"persist-scan {pold*1000:.1f}us -> {pnew*1000:.1f}us") + + +if __name__ == "__main__": + test_parity_sanitize_cursor() + test_parity_token_memo() + test_parity_persist_bounded_scan() + bench() + print("ALL PARITY CHECKS PASSED")