diff --git a/api/models.py b/api/models.py index 15552a7f..49a9ecbc 100644 --- a/api/models.py +++ b/api/models.py @@ -8259,6 +8259,66 @@ def merge_session_messages_append_only( """ sidecar_messages = list(sidecar_messages or []) state_messages = list(state_messages or []) + # Per-invocation cache keyed by message identity. Sidecar/state message objects + # are retained for this call, and this function does not mutate key-defining + # fields before each helper call. + _MESSAGE_CACHE_MISSING = object() + _cached_msg_prepared: dict[int, dict[str, object]] = {} + _cached_msg_keys: dict[tuple[int, str], object] = {} + + _message_key_helpers = { + "merge": _session_message_merge_key, + "dedup": _session_message_dedup_key, + "content": _session_message_content_key, + "visible": _session_message_visible_key, + } + + def _cached_message_key(msg, kind): + if not isinstance(msg, dict): + return _message_key_helpers[kind](msg) + + cache_key = (id(msg), kind) + value = _cached_msg_keys.get(cache_key, _MESSAGE_CACHE_MISSING) + if value is not _MESSAGE_CACHE_MISSING: + return value + + helper = _message_key_helpers[kind] + msg_cache_key = id(msg) + prepared_msg = _cached_msg_prepared.get(msg_cache_key) + + if kind in {"merge", "dedup"}: + if prepared_msg is None: + value = helper(msg) + # If this is a legacy message key, keep the already-stringified + # content payload for downstream helper calls. + if isinstance(value, tuple) and value and value[0] == "legacy": + prepared_msg = dict(msg) + prepared_msg["content"] = value[2] + _cached_msg_prepared[msg_cache_key] = prepared_msg + else: + value = helper(prepared_msg) + _cached_msg_keys[cache_key] = value + return value + + if prepared_msg is None: + # For non-ID messages this is the canonical merge path. + merge_key = _cached_message_key(msg, "merge") + prepared_msg = _cached_msg_prepared.get(msg_cache_key) + if prepared_msg is None: + prepared_msg = dict(msg) + prepared_msg["content"] = ( + merge_key[2] + if isinstance(merge_key, tuple) + and len(merge_key) > 2 + and merge_key[0] == "legacy" + else str(msg.get("content") or "") + ) + _cached_msg_prepared[msg_cache_key] = prepared_msg + + value = helper(prepared_msg) + _cached_msg_keys[cache_key] = value + return value + watermark_timestamp = _message_timestamp_as_float({"timestamp": truncation_watermark}) if not state_messages: return sidecar_messages @@ -8319,7 +8379,7 @@ def merge_session_messages_append_only( seen_messages = {} deduped = [] for msg in filtered: - key = _session_message_dedup_key(msg) + key = _cached_message_key(msg, "dedup") if key not in seen: seen.add(key) seen_messages[key] = msg @@ -8345,19 +8405,20 @@ def merge_session_messages_append_only( def _remember_merged_message(message): if not isinstance(message, dict): return - merged_by_message_key.setdefault(_session_message_merge_key(message), message) - merged_by_dedup_key.setdefault(_session_message_dedup_key(message), message) - merged_by_visible_key.setdefault(_session_message_visible_key(message), message) + merged_by_message_key.setdefault(_cached_message_key(message, "merge"), message) + merged_by_dedup_key.setdefault(_cached_message_key(message, "dedup"), message) + merged_by_visible_key.setdefault(_cached_message_key(message, "visible"), message) for msg in sidecar_messages: timestamp = _message_timestamp_as_float(msg) if timestamp is not None: max_sidecar_timestamp = timestamp if max_sidecar_timestamp is None else max(max_sidecar_timestamp, timestamp) - key = _session_message_merge_key(msg) + key = _cached_message_key(msg, "merge") seen_message_keys.add(key) - seen_dedup_keys.add(_session_message_dedup_key(msg)) - seen_content_keys.add(_session_message_content_key(msg)) - visible_key = _session_message_visible_key(msg) + seen_dedup_keys.add(_cached_message_key(msg, "dedup")) + content_key = _cached_message_key(msg, "content") + seen_content_keys.add(content_key) + visible_key = _cached_message_key(msg, "visible") seen_visible_keys.add(visible_key) sidecar_visible_keys.add(visible_key) sidecar_visible_counts[visible_key] = sidecar_visible_counts.get(visible_key, 0) + 1 @@ -8384,8 +8445,10 @@ def merge_session_messages_append_only( ) for msg in state_messages: timestamp = _message_timestamp_as_float(msg) - key = _session_message_merge_key(msg) - visible_key = _session_message_visible_key(msg) + key = _cached_message_key(msg, "merge") + dedup_key = _cached_message_key(msg, "dedup") + visible_key = _cached_message_key(msg, "visible") + content_key = _cached_message_key(msg, "content") replays_sidecar_prefix = False replay_target = None if state_replay_idx < len(sidecar_visible_sequence): @@ -8409,7 +8472,7 @@ def merge_session_messages_append_only( ) # Record dedup key so later duplicates of this replayed message # are caught by the dedup guard (#3346). - seen_dedup_keys.add(_session_message_dedup_key(msg)) + seen_dedup_keys.add(dedup_key) continue # Skip rows ABOVE the watermark only while the sidecar has NOT advanced # past the watermark. Because Session.save() no longer auto-clears the @@ -8468,7 +8531,7 @@ def merge_session_messages_append_only( and timestamp is not None and timestamp < watermark_timestamp and key not in seen_message_keys - and _session_message_content_key(msg) not in seen_content_keys + and content_key not in seen_content_keys ): continue # Same-second edit: if timestamp equals the watermark and the message @@ -8486,7 +8549,7 @@ def merge_session_messages_append_only( and timestamp is not None and timestamp == watermark_timestamp and key not in seen_message_keys - and _session_message_content_key(msg) not in seen_content_keys + and content_key not in seen_content_keys and str(msg.get("role", "")).lower() == "user" ): continue @@ -8495,7 +8558,6 @@ def merge_session_messages_append_only( # sub-second messages with the same second-level merge key are not # collapsed. The merge key truncates to seconds; the dedup key does # not. - dedup_key = _session_message_dedup_key(msg) if dedup_key in seen_dedup_keys: _merge_session_display_metadata(merged_by_dedup_key.get(dedup_key), msg) continue @@ -8563,7 +8625,7 @@ def merge_session_messages_append_only( and timestamp == watermark_timestamp and checkpoint_consumed and str(msg.get("role", "")).lower() != "user" - and _session_message_content_key(msg) not in seen_content_keys + and content_key not in seen_content_keys ): pass # fall through to append below else: @@ -8574,7 +8636,7 @@ def merge_session_messages_append_only( # differ), preserve it — distinct tool_calls must not be collapsed. _tc = msg.get("tool_calls") if _tc: - _ck = _session_message_content_key(msg) + _ck = content_key if _ck in seen_content_keys and dedup_key not in seen_dedup_keys: # Different tool_calls from sidecar — preserve, but keep # the row in timestamp order. Falling through to the @@ -8583,7 +8645,7 @@ def merge_session_messages_append_only( if _insert_state_message_chronologically(merged_messages, msg): seen_message_keys.add(key) seen_dedup_keys.add(dedup_key) - seen_content_keys.add(_session_message_content_key(msg)) + seen_content_keys.add(content_key) seen_visible_keys.add(visible_key) _remember_merged_message(msg) continue @@ -8591,11 +8653,11 @@ def merge_session_messages_append_only( _merge_session_display_metadata(merged_by_message_key.get(key), msg) continue else: - if msg.get("role") == "user" and _session_message_content_key(msg) not in seen_content_keys: + if msg.get("role") == "user" and content_key not in seen_content_keys: if _insert_state_message_chronologically(merged_messages, msg): seen_message_keys.add(key) seen_dedup_keys.add(dedup_key) - seen_content_keys.add(_session_message_content_key(msg)) + seen_content_keys.add(content_key) seen_visible_keys.add(visible_key) _remember_merged_message(msg) continue @@ -8603,7 +8665,7 @@ def merge_session_messages_append_only( continue seen_message_keys.add(key) seen_dedup_keys.add(dedup_key) - seen_content_keys.add(_session_message_content_key(msg)) + seen_content_keys.add(content_key) seen_visible_keys.add(visible_key) merged_messages.append(msg) _remember_merged_message(msg) diff --git a/tests/test_merge_append_only_normalization_cache.py b/tests/test_merge_append_only_normalization_cache.py new file mode 100644 index 00000000..1f0a0f25 --- /dev/null +++ b/tests/test_merge_append_only_normalization_cache.py @@ -0,0 +1,217 @@ +from collections import defaultdict + +import api.models as models + + +class _CountingList(list): + def __init__(self, *args, str_counter): + super().__init__(*args) + self._str_counter = str_counter + + def __str__(self): + self._str_counter["count"] += 1 + return super().__str__() + + +class _ExplodingList(list): + def __init__(self, *args, str_counter): + super().__init__(*args) + self._str_counter = str_counter + + def __str__(self): + self._str_counter["count"] += 1 + raise RuntimeError("content __str__ must not run") + + +def _install_key_wrappers(monkeypatch): + call_counts: dict[str, dict[int, int]] = { + "merge": defaultdict(int), + "dedup": defaultdict(int), + "content": defaultdict(int), + "visible": defaultdict(int), + } + + def _install(name): + original = getattr(models, f"_session_message_{name}_key") + + def _wrapped(message, *args, **kwargs): + if isinstance(message, dict): + call_counts[name][id(message)] += 1 + return original(message, *args, **kwargs) + + monkeypatch.setattr(models, f"_session_message_{name}_key", _wrapped) + + for name in call_counts: + _install(name) + + return call_counts + + +def test_merge_append_only_caches_canonical_keys_and_preserves_identity(monkeypatch): + """Verify exact key-call fanout and output identity for shared row objects.""" + str_calls = {"count": 0} + repeated_user_content = _CountingList(["hello", "world"], str_counter=str_calls) + repeated_assistant_content = _CountingList(["reply", "with", "spacing"], str_counter=str_calls) + + repeated_user = { + "role": "user", + "content": repeated_user_content, + "timestamp": 1000, + } + repeated_assistant = { + "role": "assistant", + "content": repeated_assistant_content, + "timestamp": 1001, + } + sidecar_messages = [repeated_user, repeated_user, repeated_user, repeated_assistant] + state_messages = [repeated_assistant, repeated_user] + + call_counts = _install_key_wrappers(monkeypatch) + merged = models.merge_session_messages_append_only(sidecar_messages, state_messages) + + assert [id(msg) for msg in merged] == [ + id(repeated_user), + id(repeated_user), + id(repeated_user), + id(repeated_assistant), + ] + assert merged == [repeated_user, repeated_user, repeated_user, repeated_assistant] + + assert str_calls["count"] == 2 + assert dict(call_counts["merge"]) == { + id(repeated_user): 1, + id(repeated_assistant): 1, + } + assert len(dict(call_counts["dedup"])) == 2 + assert len(dict(call_counts["content"])) == 2 + assert len(dict(call_counts["visible"])) == 2 + assert set(dict(call_counts["dedup"]).values()) == {1} + assert set(dict(call_counts["content"]).values()) == {1} + assert set(dict(call_counts["visible"]).values()) == {1} + + +def test_merge_append_only_helpers_do_not_see_mutated_content(monkeypatch): + original_content = _CountingList(["streamed", "content"], str_counter={"count": 0}) + sidecar_messages = [ + { + "role": "user", + "content": original_content, + "timestamp": 1000, + } + ] + state_messages = [ + { + "role": "assistant", + "content": "ack", + "timestamp": 1001, + } + ] + + observed = {"helpers_invoked": False} + call_counts: dict[str, dict[int, int]] = { + "merge": defaultdict(int), + "dedup": defaultdict(int), + "content": defaultdict(int), + "visible": defaultdict(int), + } + + def _install(name, *, assert_content_content=False): + original = getattr(models, f"_session_message_{name}_key") + + def _wrapped(message, *args, **kwargs): + if isinstance(message, dict): + call_counts[name][id(message)] += 1 + if assert_content_content and isinstance(message, dict): + # Verify the aliased sidecar dict still owns the original object + # during helper execution. + assert sidecar_messages[0]["content"] is original_content + observed["helpers_invoked"] = True + return original(message, *args, **kwargs) + + monkeypatch.setattr(models, f"_session_message_{name}_key", _wrapped) + + _install("merge") + _install("dedup") + _install("visible") + _install("content", assert_content_content=True) + + models.merge_session_messages_append_only(sidecar_messages, state_messages) + + assert observed["helpers_invoked"] is True + assert len(dict(call_counts["content"])) == 2 + assert set(dict(call_counts["content"]).values()) == {1} + + +def test_merge_append_only_state_only_stable_id_does_not_stringify_content(monkeypatch): + """State-only merge keeps stable-ID rows on the helper fast path.""" + str_calls = {"count": 0} + message = { + "id": "stable-id", + "role": "assistant", + "content": _ExplodingList(["danger", "zone"], str_counter=str_calls), + "timestamp": 2000, + } + + call_counts = _install_key_wrappers(monkeypatch) + merged = models.merge_session_messages_append_only([], [message]) + + assert merged == [message] + assert merged[0] is message + assert str_calls["count"] == 0 + assert dict(call_counts["merge"]) == {} + assert dict(call_counts["dedup"]) == { + id(message): 1, + } + assert dict(call_counts["content"]) == {} + assert dict(call_counts["visible"]) == {} + + +def test_merge_append_only_distinct_equal_dicts_remain_distinct_inputs(monkeypatch): + """Two equal dict rows still pass through as separate inputs.""" + msg_1 = { + "role": "user", + "content": "same", + "timestamp": 1000, + } + msg_2 = { + "role": "user", + "content": "same", + "timestamp": 1000, + } + + call_counts = _install_key_wrappers(monkeypatch) + merged = models.merge_session_messages_append_only([], [msg_1, msg_2]) + + assert merged == [msg_1] + assert merged[0] is msg_1 + assert dict(call_counts["dedup"]) == { + id(msg_1): 1, + id(msg_2): 1, + } + assert dict(call_counts["merge"]) == {} + + +def test_merge_append_only_recomputes_after_mutation(monkeypatch): + """Per-call key caches reset; mutation between calls changes work input.""" + str_calls = {"count": 0} + state_message = { + "role": "assistant", + "content": _CountingList(["first"], str_counter=str_calls), + "timestamp": 3000, + } + + call_counts = _install_key_wrappers(monkeypatch) + first = models.merge_session_messages_append_only([], [state_message]) + state_message["content"].append("second") + second = models.merge_session_messages_append_only([], [state_message]) + + assert first == [state_message] + assert second == [state_message] + assert first is not second + assert str_calls["count"] == 2 + assert dict(call_counts["dedup"]) == { + id(state_message): 2, + } + assert dict(call_counts["merge"]) == {} + assert dict(call_counts["content"]) == {} + assert dict(call_counts["visible"]) == {}