Merge pull request #6193 from starship-s/fix/session-truncate-linear-alignment

perf(session): linearize context truncation alignment

api/session_ops.py

@@ -8,6 +8,7 @@ the hermes-agent repo.
 from __future__ import annotations
 import json
 import logging
+from bisect import bisect_left
 from typing import Any
 
 from api.config import LOCK, _get_session_agent_lock
@@ -144,44 +145,201 @@ def truncate_context_for_display_keep(
             tool_calls_sig,
         )
 
-    def _first_match_from(message: Any, start_idx: int) -> tuple[int | None, int | None]:
+    # Materialize signatures once.  The matcher deliberately keeps the original
+    # rows in ``ctx``; these records are only an alignment index. A signature
+    # failure is deferred because the old matcher may return before reaching it.
+    context_records = []
+    deferred_signature_positions: list[int] = []
+    for idx, row in enumerate(ctx):
+        try:
+            row_signature = _row_signature(row)
+        except Exception:
+            row_signature = None
+            deferred_signature_positions.append(idx)
+        context_records.append((row, row_signature))
+    message_signatures = [_row_signature(message) for message in msgs]
+    id_positions: dict[Any, list[int]] = {}
+    signature_positions: dict[tuple[str, ...], list[int]] = {}
+    signature_no_id_positions: dict[tuple[str, ...], list[int]] = {}
+    signature_no_timestamp_positions: dict[tuple[str, ...], list[int]] = {}
+    signature_no_id_no_timestamp_positions: dict[tuple[str, ...], list[int]] = {}
+    signature_timestamp_positions: dict[
+        tuple[tuple[str, ...], Any], list[int]
+    ] = {}
+    signature_timestamp_no_id_positions: dict[
+        tuple[tuple[str, ...], Any], list[int]
+    ] = {}
+    unsafe_id_positions: list[int] = []
+    unsafe_timestamp_positions: dict[tuple[str, ...], list[int]] = {}
+    unsafe_timestamp_no_id_positions: dict[tuple[str, ...], list[int]] = {}
+
+    def _safe_raw_value(value: Any) -> bool:
+        # Keep dict lookup semantics aligned with the old explicit ``==`` scan:
+        # only ordinary built-in metadata values may use the raw-value indexes.
+        # In particular, custom objects and non-reflexive NaN values can make a
+        # dict find a key that the old equality check rejected.
+        if value is None:
+            return True
+        if type(value) not in (str, int, float):
+            return False
+        try:
+            hash(value)
+            return value == value
+        except Exception:
+            return False
+
+    for idx, (context_row, context_sig) in enumerate(context_records):
+        if context_sig is not None:
+            signature_positions.setdefault(context_sig, []).append(idx)
+        if not isinstance(context_row, dict):
+            continue
+        context_id = context_row.get('id')
+        context_ts = context_row.get('timestamp')
+        if context_sig is not None:
+            if context_id is None:
+                signature_no_id_positions.setdefault(context_sig, []).append(idx)
+            if context_ts is None:
+                signature_no_timestamp_positions.setdefault(context_sig, []).append(idx)
+            if context_id is None and context_ts is None:
+                signature_no_id_no_timestamp_positions.setdefault(
+                    context_sig, []
+                ).append(idx)
+        if context_id is not None and _safe_raw_value(context_id):
+            id_positions.setdefault(context_id, []).append(idx)
+        elif context_id is not None:
+            unsafe_id_positions.append(idx)
+        if (
+            context_sig is not None
+            and context_ts is not None
+            and _safe_raw_value(context_ts)
+        ):
+            timestamp_key = (context_sig, context_ts)
+            signature_timestamp_positions.setdefault(timestamp_key, []).append(idx)
+            if context_id is None:
+                signature_timestamp_no_id_positions.setdefault(
+                    timestamp_key, []
+                ).append(idx)
+        elif context_sig is not None and context_ts is not None:
+            unsafe_timestamp_positions.setdefault(context_sig, []).append(idx)
+            if context_id is None:
+                unsafe_timestamp_no_id_positions.setdefault(
+                    context_sig, []
+                ).append(idx)
+
+    def _first_at_or_after(positions: list[int] | None, start_idx: int) -> int | None:
+        if not positions:
+            return None
+        offset = bisect_left(positions, start_idx)
+        return positions[offset] if offset < len(positions) else None
+
+    def _lazy_first_match_from(
+        message: Any,
+        start_idx: int,
+    ) -> tuple[int | None, int | None]:
+        """Match exactly as the original ordered scan did."""
         msg_sig = _row_signature(message)
         if msg_sig is None:
             return None, None
-        msg_id = message.get('id')
-        msg_ts = message.get('timestamp')
         weak_matches: list[int] = []
         for idx in range(start_idx, len(ctx)):
             context_row = ctx[idx]
             context_sig = _row_signature(context_row)
-            if context_sig is None:
+            if context_sig is None or not isinstance(context_row, dict):
                 continue
-
             context_id = context_row.get('id')
+            msg_id = message.get('id')
             if context_id is not None and msg_id is not None:
                 if context_id == msg_id:
                     return idx, None
                 continue
-
             if context_sig != msg_sig:
                 continue
-
             context_ts = context_row.get('timestamp')
+            msg_ts = message.get('timestamp')
             if context_ts is not None and msg_ts is not None:
                 if context_ts == msg_ts:
                     return idx, None
                 continue
-
             weak_matches.append(idx)
             if len(weak_matches) > 1:
                 return None, weak_matches[0]
         return (weak_matches[0], None) if len(weak_matches) == 1 else (None, None)
 
+    def _first_match_from(
+        message_idx: int,
+        message: Any,
+        start_idx: int,
+    ) -> tuple[int | None, int | None]:
+        msg_sig = message_signatures[message_idx]
+        if msg_sig is None:
+            return None, None
+        msg_id = message.get('id')
+        msg_ts = message.get('timestamp')
+        deferred_reachable = _first_at_or_after(
+            deferred_signature_positions, start_idx
+        ) is not None
+        unsafe_id_reachable = (
+            msg_id is not None
+            and _first_at_or_after(unsafe_id_positions, start_idx) is not None
+        )
+        unsafe_timestamp_candidates = (
+            unsafe_timestamp_no_id_positions.get(msg_sig, [])
+            if msg_id is not None
+            else unsafe_timestamp_positions.get(msg_sig, [])
+        )
+        unsafe_timestamp_reachable = (
+            msg_ts is not None
+            and _first_at_or_after(unsafe_timestamp_candidates, start_idx) is not None
+        )
+        if not _safe_raw_value(msg_id) or not _safe_raw_value(msg_ts):
+            return _lazy_first_match_from(message, start_idx)
+        if deferred_reachable or unsafe_id_reachable or unsafe_timestamp_reachable:
+            return _lazy_first_match_from(message, start_idx)
+
+        exact_positions: list[int] = []
+        if msg_id is not None:
+            id_idx = _first_at_or_after(id_positions.get(msg_id), start_idx)
+            if id_idx is not None:
+                exact_positions.append(id_idx)
+        if msg_ts is not None:
+            timestamp_key = (msg_sig, msg_ts)
+            if msg_id is not None:
+                timestamp_positions = signature_timestamp_no_id_positions.get(
+                    timestamp_key, []
+                )
+            else:
+                timestamp_positions = signature_timestamp_positions.get(
+                    timestamp_key, []
+                )
+            timestamp_idx = _first_at_or_after(timestamp_positions, start_idx)
+            if timestamp_idx is not None:
+                exact_positions.append(timestamp_idx)
+        exact_idx = min(exact_positions, default=None)
+
+        if msg_id is not None and msg_ts is not None:
+            weak_candidates = signature_no_id_no_timestamp_positions.get(msg_sig)
+        elif msg_id is not None:
+            weak_candidates = signature_no_id_positions.get(msg_sig)
+        elif msg_ts is not None:
+            weak_candidates = signature_no_timestamp_positions.get(msg_sig)
+        else:
+            weak_candidates = signature_positions.get(msg_sig)
+        weak_start = bisect_left(weak_candidates, start_idx) if weak_candidates else 0
+        weak_positions = weak_candidates[weak_start:weak_start + 2] if weak_candidates else []
+        second_weak_idx = weak_positions[1] if len(weak_positions) > 1 else None
+        if second_weak_idx is not None and (
+            exact_idx is None or second_weak_idx < exact_idx
+        ):
+            return None, weak_positions[0]
+        if exact_idx is not None:
+            return exact_idx, None
+        return (weak_positions[0], None) if len(weak_positions) == 1 else (None, None)
+
     matches = [None] * len(msgs)
     ambiguous_matches = [None] * len(msgs)
     next_ctx_idx = 0
     for msg_idx, message in enumerate(msgs):
-        match_idx, ambiguous_idx = _first_match_from(message, next_ctx_idx)
+        match_idx, ambiguous_idx = _first_match_from(msg_idx, message, next_ctx_idx)
         matches[msg_idx] = match_idx
         ambiguous_matches[msg_idx] = ambiguous_idx
         if match_idx is not None:

tests/test_session_truncate_alignment_performance.py

@@ -0,0 +1,454 @@
+"""Regression coverage for session context/display alignment."""
+
+import copy
+import random
+
+import api.session_ops as session_ops
+import pytest
+
+
+def _old_matcher(context, display, keep):
+    """Test-local copy of the pre-indexed matcher for differential checks."""
+    if keep <= 0:
+        return []
+    ctx = context if isinstance(context, list) else []
+    display = display if isinstance(display, list) else []
+    if not ctx or not display:
+        return []
+    if len(ctx) == len(display):
+        return ctx[:keep]
+
+    def signature(row):
+        if not isinstance(row, dict):
+            return None
+        tool_calls = row.get("tool_calls")
+        return (
+            str(row.get("role") or ""),
+            str(row.get("content") or ""),
+            str(row.get("tool_call_id") or ""),
+            str(row.get("tool_use_id") or ""),
+            str(row.get("tool_name") or row.get("name") or ""),
+            session_ops.json.dumps(tool_calls, sort_keys=True, default=str)
+            if tool_calls else "",
+        )
+
+    def first_match(message, start):
+        msg_sig = signature(message)
+        if msg_sig is None:
+            return None, None
+        weak = []
+        for idx in range(start, len(ctx)):
+            row = ctx[idx]
+            row_sig = signature(row)
+            if row_sig is None:
+                continue
+            row_id = row.get("id")
+            msg_id = message.get("id")
+            if row_id is not None and msg_id is not None:
+                if row_id == msg_id:
+                    return idx, None
+                continue
+            if row_sig != msg_sig:
+                continue
+            row_ts = row.get("timestamp")
+            msg_ts = message.get("timestamp")
+            if row_ts is not None and msg_ts is not None:
+                if row_ts == msg_ts:
+                    return idx, None
+                continue
+            weak.append(idx)
+            if len(weak) > 1:
+                return None, weak[0]
+        return (weak[0], None) if len(weak) == 1 else (None, None)
+
+    matches = []
+    ambiguous = []
+    start = 0
+    for message in display:
+        match, ambiguous_match = first_match(message, start)
+        matches.append(match)
+        ambiguous.append(ambiguous_match)
+        if match is not None:
+            start = match + 1
+    if keep < len(display):
+        last = matches[keep - 1] if keep else None
+        first = matches[keep]
+        if first is not None:
+            if (
+                last is not None
+                and isinstance(display[keep - 1], dict)
+                and display[keep - 1].get("role") == "user"
+            ):
+                return ctx[:last + 1]
+            return ctx[:first]
+        if last is not None:
+            if (
+                ambiguous[keep] is not None
+                and isinstance(display[keep - 1], dict)
+                and display[keep - 1].get("role") != "user"
+            ):
+                return ctx[:ambiguous[keep]]
+            return ctx[:last + 1]
+        if len(ctx) < len(display):
+            for i in range(keep - 1, -1, -1):
+                resolved = matches[i] if matches[i] is not None else ambiguous[i]
+                if resolved is not None:
+                    return ctx[:resolved + 1]
+    prefix_len = max(0, len(ctx) - len(display))
+    return ctx[:prefix_len] + ctx[prefix_len:prefix_len + keep]
+
+
+def test_alignment_constructs_each_signature_once(monkeypatch):
+    display = [
+        {"role": "assistant", "content": f"display-{i}", "tool_calls": [{"id": i}], "id": f"d-{i}"}
+        for i in range(30)
+    ]
+    context = [
+        {"role": "assistant", "content": f"context-{i}", "tool_calls": [{"id": i}], "id": f"c-{i}"}
+        for i in range(15)
+    ]
+    calls = 0
+    original = session_ops.json.dumps
+
+    def counted(*args, **kwargs):
+        nonlocal calls
+        calls += 1
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(session_ops.json, "dumps", counted)
+    session_ops.truncate_context_for_display_keep(context, display, 10)
+    assert calls <= len(display) + len(context)
+
+
+def test_indexed_alignment_matches_reference_for_ordering_hazards():
+    cases = [
+        # One weak candidate before an exact id match: exact wins.
+        (
+            [{"role": "assistant", "content": "x"}, {"role": "assistant", "content": "x", "id": "exact"}, {"role": "assistant", "content": "context-tail"}],
+            [{"role": "assistant", "content": "x", "id": "exact"}, {"role": "assistant", "content": "tail"}],
+        ),
+        # Two weak candidates before a later exact match: remain ambiguous.
+        (
+            [{"role": "assistant", "content": "x"}, {"role": "assistant", "content": "x"}, {"role": "assistant", "content": "x", "id": "exact"}, {"role": "assistant", "content": "context-tail"}],
+            [{"role": "assistant", "content": "x", "id": "exact"}, {"role": "assistant", "content": "tail"}],
+        ),
+        # The same ordering rule using timestamps for the exact candidate.
+        (
+            [{"role": "assistant", "content": "x"}, {"role": "assistant", "content": "x", "timestamp": 2}, {"role": "assistant", "content": "context-tail"}],
+            [{"role": "assistant", "content": "x", "timestamp": 2}, {"role": "assistant", "content": "tail"}],
+        ),
+    ]
+    for context, display in cases:
+        expected = _old_matcher(copy.deepcopy(context), copy.deepcopy(display), 1)
+        actual = session_ops.truncate_context_for_display_keep(context, display, 1)
+        assert actual == expected
+
+
+def test_alignment_accepts_unhashable_metadata_and_malformed_rows():
+    context = [None, {"role": "assistant", "content": "x", "id": ["same"], "timestamp": {"n": 1}}, {"role": "assistant", "content": "context-tail"}]
+    display = [
+        {"role": "assistant", "content": "x", "id": ["same"], "timestamp": {"n": 1}},
+        {"role": "assistant", "content": "tail"},
+    ]
+    assert session_ops.truncate_context_for_display_keep(context, display, 1) == context[:2]
+
+
+class _HashableValue:
+    def __init__(self, value):
+        self.value = value
+
+    def __eq__(self, other):
+        return getattr(other, "value", object()) == self.value
+
+    def __hash__(self):
+        return hash(self.value)
+
+
+class _NonReflexiveHashableValue(_HashableValue):
+    def __eq__(self, other):
+        return False
+
+    __hash__ = _HashableValue.__hash__
+
+
+class _UnhashableValue(_HashableValue):
+    __hash__ = None
+
+
+class _UnhashableEqualValue:
+    def __init__(self, value):
+        self.value = value
+
+    def __eq__(self, other):
+        return self.value == other
+
+    __hash__ = None
+
+
+def test_alignment_preserves_cross_hashability_equality():
+    context = [
+        {"role": "assistant", "content": "id", "id": _UnhashableValue("same")},
+        {
+            "role": "assistant",
+            "content": "timestamp",
+            "timestamp": _UnhashableValue("same-time"),
+        },
+        {"role": "assistant", "content": "tail"},
+    ]
+    display = [
+        {"role": "assistant", "content": "id", "id": _HashableValue("same")},
+        {
+            "role": "assistant",
+            "content": "timestamp",
+            "timestamp": _HashableValue("same-time"),
+        },
+        {"role": "assistant", "content": "tail"},
+    ]
+    expected = _old_matcher(copy.deepcopy(context), copy.deepcopy(display), 2)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 2)
+    assert actual == expected == context[:2]
+
+
+def test_alignment_matches_old_matcher_for_non_reflexive_id():
+    nan = float("nan")
+    context = [
+        {"role": "assistant", "content": "x", "id": nan},
+    ]
+    display = [
+        {"role": "assistant", "content": "x", "id": nan},
+        {"role": "assistant", "content": "x"},
+    ]
+
+    expected = _old_matcher(context, display, 1)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 1)
+    assert actual == expected == []
+
+
+def test_alignment_matches_unsafe_context_id_against_safe_display_id():
+    context = [
+        {"role": "system", "content": "prefix"},
+        {"role": "assistant", "content": "id-target", "id": True},
+        {"role": "assistant", "content": "tail"},
+    ]
+    display = [
+        {"role": "assistant", "content": "id-target", "id": 1},
+        {"role": "assistant", "content": "tail"},
+    ]
+
+    expected = _old_matcher(copy.deepcopy(context), copy.deepcopy(display), 1)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 1)
+    assert actual == expected == context[:2]
+
+
+def test_alignment_matches_unsafe_context_timestamp_against_safe_display_timestamp():
+    context = [
+        {"role": "system", "content": "prefix"},
+        {
+            "role": "assistant",
+            "content": "timestamp-target",
+            "timestamp": _UnhashableEqualValue(7),
+        },
+        {"role": "assistant", "content": "tail"},
+    ]
+    display = [
+        {"role": "assistant", "content": "timestamp-target", "timestamp": 7},
+        {"role": "assistant", "content": "tail"},
+    ]
+
+    expected = _old_matcher(copy.deepcopy(context), copy.deepcopy(display), 1)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 1)
+    assert actual == expected == context[:2]
+
+
+def test_alignment_matches_old_matcher_for_non_reflexive_timestamp():
+    nan = float("nan")
+    context = [
+        {"role": "assistant", "content": "x", "timestamp": nan},
+    ]
+    display = [
+        {"role": "assistant", "content": "x", "timestamp": nan},
+        {"role": "assistant", "content": "x"},
+    ]
+
+    expected = _old_matcher(context, display, 1)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 1)
+    assert actual == expected == []
+
+
+def test_alignment_matches_old_matcher_for_non_reflexive_hashable_object():
+    value = _NonReflexiveHashableValue("same")
+    context = [{"role": "assistant", "content": "x", "id": value}]
+    display = [
+        {"role": "assistant", "content": "x", "id": value},
+        {"role": "assistant", "content": "x"},
+    ]
+
+    expected = _old_matcher(context, display, 1)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 1)
+    assert actual == expected == []
+
+
+def test_shorter_context_reverse_fallback_prefers_ambiguous_kept_boundary():
+    exact_row = {"role": "assistant", "content": "exact", "id": "exact"}
+    context = [
+        exact_row,
+        {"role": "assistant", "content": "x"},
+        {"role": "assistant", "content": "x"},
+    ]
+    display = [
+        {"role": "assistant", "content": "exact", "id": "exact"},
+        {"role": "assistant", "content": "x"},
+        {"role": "assistant", "content": "unmatched boundary"},
+        {"role": "assistant", "content": "extra unmatched"},
+    ]
+
+    expected = _old_matcher(copy.deepcopy(context), copy.deepcopy(display), 2)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 2)
+    assert actual == expected == context[:2]
+
+
+def test_alignment_differential_randomized_against_origin_matcher():
+    rng = random.Random(5096)
+    roles = ["user", "assistant", "tool", None]
+    collision_values = [True, False, 0, 0.0, 1, 1.0]
+    for _ in range(500):
+        context_len = rng.randrange(0, 30)
+        display_len = rng.randrange(0, 30)
+        def make_row(index):
+            if rng.random() < 0.15:
+                return rng.choice([None, "malformed", 17])
+            row = {
+                "role": rng.choice(roles),
+                "content": rng.choice(["same", "other", ""]),
+            }
+            if rng.random() < 0.55:
+                row["id"] = rng.choice(["a", "b", None, *collision_values])
+            if rng.random() < 0.55:
+                row["timestamp"] = rng.choice([1, 2, None, *collision_values])
+            if rng.random() < 0.35:
+                row["tool_calls"] = [{"id": index % 2}]
+            return row
+
+        context = [make_row(i) for i in range(context_len)]
+        display = [make_row(i + 20) for i in range(display_len)]
+        keep = rng.randrange(-1, display_len + 2)
+        expected = _old_matcher(copy.deepcopy(context), copy.deepcopy(display), keep)
+        actual = session_ops.truncate_context_for_display_keep(context, display, keep)
+        assert actual == expected, (context, display, keep)
+
+
+class _CountingDict(dict):
+    def __init__(self, counter, **kwargs):
+        super().__init__(**kwargs)
+        self._counter = counter
+
+    def get(self, key, default=None):
+        if key == "id":
+            self._counter["id_gets"] += 1
+        return super().get(key, default)
+
+
+def test_alignment_prefilters_timestamp_candidates_once():
+    counter = {"id_gets": 0}
+    context = [
+        _CountingDict(
+            counter, role="assistant", content="same", timestamp=7, id=f"ctx-{i}"
+        )
+        for i in range(500)
+    ]
+    display = [
+        {"role": "assistant", "content": "same", "timestamp": 7, "id": f"msg-{i}"}
+        for i in range(250)
+    ]
+    session_ops.truncate_context_for_display_keep(context, display, 125)
+    assert counter["id_gets"] <= len(context) + 2
+
+
+def test_alignment_preserves_reached_raising_context_signature():
+    context = [
+        {
+            "role": "assistant",
+            "content": "reached",
+            "tool_calls": [{1: "integer key", "string key": "mixed keys"}],
+        },
+        {"role": "assistant", "content": "context-tail", "id": "tail"},
+    ]
+    display = [
+        {"role": "assistant", "content": "display-head"},
+        {"role": "assistant", "content": "display-tail"},
+        {"role": "assistant", "content": "display-extra"},
+    ]
+
+    with pytest.raises(TypeError):
+        _old_matcher(context, display, 1)
+    with pytest.raises(TypeError):
+        session_ops.truncate_context_for_display_keep(context, display, 1)
+
+
+def test_alignment_defers_unreachable_raising_context_signature():
+    context = [
+        {"role": "assistant", "content": "first", "id": "first"},
+        {"role": "assistant", "content": "second", "id": "second"},
+        {
+            "role": "assistant",
+            "content": "unreachable",
+            "tool_calls": [{1: "integer key", "string key": "mixed keys"}],
+        },
+    ]
+    display = [
+        {"role": "assistant", "content": "first", "id": "first"},
+        {"role": "assistant", "content": "second", "id": "second"},
+    ]
+
+    expected = _old_matcher(context, display, 1)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 1)
+
+    assert expected == context[:1]
+    assert actual == expected
+
+
+class _UnsafeEquality:
+    def __eq__(self, other):
+        raise AssertionError("unsafe equality was reached")
+
+    __hash__ = object.__hash__
+
+
+def test_alignment_stops_before_unsafe_id_after_weak_ambiguity():
+    context = [
+        {"role": "assistant", "content": "same"},
+        {"role": "assistant", "content": "same"},
+        {"role": "assistant", "content": "same", "id": _UnsafeEquality()},
+    ]
+    display = [
+        {"role": "assistant", "content": "same", "id": "target"},
+        {"role": "assistant", "content": "tail"},
+    ]
+
+    expected = _old_matcher(context, display, 1)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 1)
+
+    assert expected == context[:2]
+    assert actual == expected
+
+
+def test_alignment_stops_before_unsafe_timestamp_after_weak_ambiguity():
+    context = [
+        {"role": "assistant", "content": "same"},
+        {"role": "assistant", "content": "same"},
+        {
+            "role": "assistant",
+            "content": "same",
+            "timestamp": _UnsafeEquality(),
+        },
+    ]
+    display = [
+        {"role": "assistant", "content": "same", "timestamp": 7},
+        {"role": "assistant", "content": "tail"},
+    ]
+
+    expected = _old_matcher(context, display, 1)
+    actual = session_ops.truncate_context_for_display_keep(context, display, 1)
+
+    assert expected == context[:2]
+    assert actual == expected