Merge pull request #6182 from starship-s/perf/session-switch-merge-key-cache
perf(sessions): cache reconciliation message keys
This commit is contained in:
2 files changed
+299
-20
No files matched your search
+82
-20
@@ -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"]) == {}
|
||||
Reference in new issue
Block a user