diff --git a/api/session_ops.py b/api/session_ops.py index 29432a31..05f0038f 100644 --- a/api/session_ops.py +++ b/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,180 @@ def truncate_context_for_display_keep( tool_calls_sig, ) - def _first_match_from(message: Any, start_idx: int) -> tuple[int | None, int | None]: - msg_sig = _row_signature(message) + # Materialize signatures once. The matcher deliberately keeps the original + # rows in ``ctx``; these records are only an alignment index. + context_records = [ + (row, _row_signature(row)) for row in ctx + ] + 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: + return value == value and hash(value) is not None + 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 _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') - 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: - continue + if not _safe_raw_value(msg_id) or not _safe_raw_value(msg_ts): + weak_matches: list[int] = [] + for idx in range(start_idx, len(ctx)): + context_row, context_sig = context_records[idx] + if context_sig is None or not isinstance(context_row, dict): + continue + context_id = context_row.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') + 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) - context_id = context_row.get('id') - if context_id is not None and msg_id is not None: - if context_id == msg_id: - return idx, None - continue + 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) + for idx in unsafe_id_positions: + if idx < start_idx: + continue + if ctx[idx].get('id') == msg_id: + exact_positions.append(idx) + break + 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, [] + ) + unsafe_timestamp_candidates = ( + unsafe_timestamp_no_id_positions.get(msg_sig, []) + ) + else: + timestamp_positions = signature_timestamp_positions.get( + timestamp_key, [] + ) + unsafe_timestamp_candidates = unsafe_timestamp_positions.get( + msg_sig, [] + ) + timestamp_idx = _first_at_or_after(timestamp_positions, start_idx) + if timestamp_idx is not None: + exact_positions.append(timestamp_idx) + for idx in unsafe_timestamp_candidates: + if idx < start_idx: + continue + if ctx[idx].get('timestamp') == msg_ts: + exact_positions.append(idx) + break + exact_idx = min(exact_positions, default=None) - if context_sig != msg_sig: - continue - - context_ts = context_row.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) + 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: diff --git a/tests/test_session_truncate_alignment_performance.py b/tests/test_session_truncate_alignment_performance.py new file mode 100644 index 00000000..6ca34613 --- /dev/null +++ b/tests/test_session_truncate_alignment_performance.py @@ -0,0 +1,293 @@ +"""Regression coverage for session context/display alignment.""" + +import copy +import random + +import api.session_ops as session_ops + + +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 + + +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_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_alignment_differential_randomized_against_origin_matcher(): + rng = random.Random(5096) + roles = ["user", "assistant", "tool", None] + for _ in range(120): + context_len = rng.randrange(0, 9) + display_len = rng.randrange(0, 9) + 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]) + if rng.random() < 0.55: + row["timestamp"] = rng.choice([1, 2, None]) + 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): + id_gets = 0 + + def get(self, key, default=None): + if key == "id": + type(self).id_gets += 1 + return super().get(key, default) + + +def test_alignment_prefilters_timestamp_candidates_once(): + _CountingDict.id_gets = 0 + context = [ + _CountingDict(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 _CountingDict.id_gets <= len(context) + 2