Merge pull request #6182 from starship-s/perf/session-switch-merge-key-cache

perf(sessions): cache reconciliation message keys
This commit is contained in:
nesquena-hermes authored and GitHub committed 2026-07-17 03:04:15 -07:00
commit f6265cc962
2 files changed
+299 -20

No files matched your search

+82 -20
View File
@@ -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)
@@ -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"]) == {}