fix: preserve lazy alignment semantics

api/session_ops.py

@@ -146,10 +146,17 @@ def truncate_context_for_display_keep(
         )
 
     # 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
-    ]
+    # 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]] = {}
@@ -225,6 +232,39 @@ def truncate_context_for_display_keep(
         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
+        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 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,
@@ -235,65 +275,45 @@ def truncate_context_for_display_keep(
             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):
-            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)
+            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)
-            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 msg_id is not None and msg_ts is not None:

tests/test_session_truncate_alignment_performance.py

@@ -4,6 +4,7 @@ import copy
 import random
 
 import api.session_ops as session_ops
+import pytest
 
 
 def _old_matcher(context, display, keep):
@@ -361,3 +362,93 @@ def test_alignment_prefilters_timestamp_candidates_once():
     ]
     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