Release v0.51.458 — O(D²) context-backfill, prevents server hang (#4314) (#4316)

Release v0.51.458 — O(D²) context-backfill, prevents server hang (#4314)

CHANGELOG.md

@@ -3,6 +3,12 @@
 
 ## [Unreleased]
 
+## [v0.51.458] — 2026-06-16 — Release PS (long sessions stop hanging)
+
+### Fixed
+
+- **Server no longer hangs on long tool-call sessions (#4314).** The transcript-backfill loop in `_merge_display_messages_after_agent_result` was O(D²·C): it re-ran `_message_identity` (`json.dumps`) for every future display row on every outer iteration and used an O(N) list-slice membership test. On sessions with many tool-call turns this pegged a CPU core, held the GIL, and stalled the `/api/sessions` sidebar poll 5–11s per cycle — the WebUI appeared hung on the current turn. The loop now precomputes display identities once and keeps an O(1) multiset mirror of the remaining context keys, making it O(D²). Behavior is unchanged: the multiset preserves the exact membership semantics of the original list-slice (including duplicate-identity turns and empty rows), verified by a differential regression test. Thanks @psanger.
+
 ## [v0.51.457] — 2026-06-16 — Release PR (gateway watcher can't brick startup)
 
 ### Fixed

api/streaming.py

@@ -4171,6 +4171,19 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex
         _has_context_only_turns = bool(_context_id_set - _display_id_set)
         if _has_context_only_turns:
             context_keys = [_message_identity(m) for m in previous_context]
+            # Precompute display keys once; avoids repeated json.dumps calls inside
+            # the inner any() loop (was O(D²·C) — see perf fix below).
+            _display_keys = [_message_identity(m) for m in previous_display]
+            # Multiset mirror of context_keys[_cursor:] kept in sync as _cursor
+            # advances. Enables O(1) membership tests in the any() check instead
+            # of an O(N) list scan, while preserving EXACT list-slice semantics:
+            # _message_identity intentionally returns duplicate keys for
+            # identical-content turns (and None for empty rows), so a plain set
+            # would drop a key still present later in the slice. A count-keyed
+            # dict (including None) matches `in context_keys[_cursor:]` exactly.
+            _remaining_ck_counts = {}
+            for _ck in context_keys:
+                _remaining_ck_counts[_ck] = _remaining_ck_counts.get(_ck, 0) + 1
             _backfilled = []
             # #3300 fix: track ONLY context rows we splice in, so the
             # visible-display backbone is never suppressed. Sharing one set
@@ -4182,7 +4195,7 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex
             _context_inserted = set()
             _cursor = 0
             for _display_idx, _dmsg in enumerate(previous_display):
-                _dkey = _message_identity(_dmsg)
+                _dkey = _display_keys[_display_idx]
                 if _dkey is not None:
                     _j = _cursor
                     while _j < len(context_keys) and context_keys[_j] != _dkey:
@@ -4194,10 +4207,20 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex
                             if _ckey is not None and _ckey not in _context_inserted and _ckey not in _display_id_set and not _is_context_compression_marker(_cmsg):
                                 _backfilled.append(copy.deepcopy(_cmsg))
                                 _context_inserted.add(_ckey)
+                        # Sync multiset: decrement keys consumed by advancing
+                        # the cursor to _j+1 (delete at zero so membership matches
+                        # the list slice exactly).
+                        for _k in range(_cursor, _j + 1):
+                            _consumed_ck = context_keys[_k]
+                            _ck_n = _remaining_ck_counts.get(_consumed_ck, 0) - 1
+                            if _ck_n <= 0:
+                                _remaining_ck_counts.pop(_consumed_ck, None)
+                            else:
+                                _remaining_ck_counts[_consumed_ck] = _ck_n
                         _cursor = _j + 1
                     elif not any(
-                        _message_identity(_future_dmsg) in context_keys[_cursor:]
-                        for _future_dmsg in previous_display[_display_idx + 1:]
+                        _display_keys[_fi] in _remaining_ck_counts
+                        for _fi in range(_display_idx + 1, len(_display_keys))
                     ):
                         for _k in range(_cursor, len(context_keys)):
                             _ckey = context_keys[_k]
@@ -4206,6 +4229,7 @@ def _merge_display_messages_after_agent_result(previous_display, previous_contex
                                 _backfilled.append(copy.deepcopy(_cmsg))
                                 _context_inserted.add(_ckey)
                         _cursor = len(context_keys)
+                        _remaining_ck_counts.clear()
                 # The display row is the visible backbone — always preserve it,
                 # in order, even when an earlier (identical-content) turn or a
                 # backfilled context row shares its timestamp-less identity.

tests/test_merge_backfill_perf_optimization.py

@@ -0,0 +1,163 @@
+"""Regression test for the O(D²·C) → O(D²) backfill optimization in
+_merge_display_messages_after_agent_result (salvaged from #4314).
+
+The optimization replaces an `in context_keys[_cursor:]` list-slice membership
+test with an O(1) count-keyed dict mirror. The subtle correctness requirement:
+_message_identity intentionally returns DUPLICATE keys for identical-content
+turns (and None for empty rows). A plain set would diverge from the original
+list-slice semantics; a multiset (count dict, including None) is exact.
+
+This test asserts the optimized merge produces byte-identical output to a
+reference implementation of the ORIGINAL list-slice semantics, over adversarial
+inputs that force duplicate identities and None keys.
+"""
+import sys
+from pathlib import Path
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+from api import streaming  # noqa: E402
+
+
+def _msg(role, text, **extra):
+    m = {"role": role, "content": text}
+    m.update(extra)
+    return m
+
+
+def test_backfill_optimization_preserves_duplicate_identity_turns():
+    """Two identical-content user turns both survive the backfill merge.
+
+    _message_identity collapses identical user text to the same key. The
+    optimization must not drop the second identical turn (a plain-set mirror
+    would). previous_display is the visible backbone; both 'Ok' user bubbles
+    plus the interleaved assistant rows must be preserved in order.
+    """
+    previous_display = [
+        _msg("user", "Ok"),
+        _msg("assistant", "first reply"),
+        _msg("user", "Ok"),
+        _msg("assistant", "second reply"),
+    ]
+    # Context has a context-only turn (never rendered) that must be backfilled,
+    # plus the same duplicate-identity 'Ok' rows.
+    previous_context = [
+        _msg("user", "Ok"),
+        _msg("assistant", "first reply"),
+        _msg("user", "context only — behind a compression marker"),
+        _msg("user", "Ok"),
+        _msg("assistant", "second reply"),
+    ]
+    result_messages = list(previous_context) + [_msg("assistant", "third reply")]
+
+    merged = streaming._merge_display_messages_after_agent_result(
+        previous_display, previous_context, result_messages, "Ok"
+    )
+
+    # At least the two backbone 'Ok' user turns survive (not collapsed to one by
+    # a buggy set-mirror that drops a still-present duplicate identity).
+    user_oks = [m for m in merged if m.get("role") == "user" and m.get("content") == "Ok"]
+    assert len(user_oks) >= 2, f"duplicate-identity user turn dropped: {merged}"
+    # The context-only turn was backfilled into the visible transcript.
+    assert any(
+        m.get("content", "").startswith("context only") for m in merged
+    ), f"context-only turn not backfilled: {merged}"
+    # Visible backbone order preserved: first reply before second reply.
+    contents = [m.get("content") for m in merged]
+    assert contents.index("first reply") < contents.index("second reply")
+
+
+def test_backfill_optimization_matches_reference_listslice_semantics():
+    """Differential check: optimized merge == reference (original) semantics
+    over adversarial inputs with duplicate identities and empty rows."""
+    import copy as _copy
+    import random
+
+    def reference_merge(previous_display, previous_context, result_messages, msg_text):
+        # Faithful re-implementation of the PRE-optimization inner loop using the
+        # original `in context_keys[_cursor:]` list-slice membership test.
+        previous_display = list(previous_display or [])
+        _partial_seen = set()
+        _deduped_rev = []
+        for m in reversed(previous_display):
+            if isinstance(m, dict) and m.get("_partial"):
+                key = streaming._message_identity(m)
+                if key is not None:
+                    if key in _partial_seen:
+                        continue
+                    _partial_seen.add(key)
+            _deduped_rev.append(m)
+        previous_display = list(reversed(_deduped_rev))
+        previous_context = list(previous_context or [])
+        result_messages = list(result_messages or [])
+        if not result_messages:
+            return previous_display
+        if previous_display and previous_context:
+            _display_id_set = {streaming._message_identity(m) for m in previous_display}
+            _context_id_set = {
+                streaming._message_identity(m)
+                for m in previous_context
+                if not streaming._is_context_compression_marker(m)
+            }
+            if bool(_context_id_set - _display_id_set):
+                context_keys = [streaming._message_identity(m) for m in previous_context]
+                _backfilled = []
+                _context_inserted = set()
+                _cursor = 0
+                for _di, _dmsg in enumerate(previous_display):
+                    _dkey = streaming._message_identity(_dmsg)
+                    if _dkey is not None:
+                        _j = _cursor
+                        while _j < len(context_keys) and context_keys[_j] != _dkey:
+                            _j += 1
+                        if _j < len(context_keys):
+                            for _k in range(_cursor, _j):
+                                _ckey = context_keys[_k]
+                                _cmsg = previous_context[_k]
+                                if _ckey is not None and _ckey not in _context_inserted and _ckey not in _display_id_set and not streaming._is_context_compression_marker(_cmsg):
+                                    _backfilled.append(_copy.deepcopy(_cmsg))
+                                    _context_inserted.add(_ckey)
+                            _cursor = _j + 1
+                        elif not any(
+                            streaming._message_identity(_f) in context_keys[_cursor:]
+                            for _f in previous_display[_di + 1:]
+                        ):
+                            for _k in range(_cursor, len(context_keys)):
+                                _ckey = context_keys[_k]
+                                _cmsg = previous_context[_k]
+                                if _ckey is not None and _ckey not in _context_inserted and _ckey not in _display_id_set and not streaming._is_context_compression_marker(_cmsg):
+                                    _backfilled.append(_copy.deepcopy(_cmsg))
+                                    _context_inserted.add(_ckey)
+                            _cursor = len(context_keys)
+                    _backfilled.append(_dmsg)
+                while _cursor < len(context_keys):
+                    _ckey = context_keys[_cursor]
+                    _cmsg = previous_context[_cursor]
+                    _cursor += 1
+                    if _ckey is not None and _ckey not in _context_inserted and _ckey not in _display_id_set and not streaming._is_context_compression_marker(_cmsg):
+                        _backfilled.append(_copy.deepcopy(_cmsg))
+                        _context_inserted.add(_ckey)
+                if len(_backfilled) > len(previous_display):
+                    previous_display = _backfilled
+        # Both share the identical tail-merge logic after backfill; compare backfill output.
+        return previous_display
+
+    rng = random.Random(2026)
+    texts = ["Ok", "Hi", "", "context only"]
+    roles = ["user", "assistant"]
+    for _ in range(2000):
+        pd = [_msg(rng.choice(roles), rng.choice(texts)) for _ in range(rng.randint(0, 5))]
+        pc = [_msg(rng.choice(roles), rng.choice(texts)) for _ in range(rng.randint(0, 6))]
+        opt = streaming._merge_display_messages_after_agent_result(
+            [dict(m) for m in pd], [dict(m) for m in pc], [dict(m) for m in pc] + [_msg("assistant", "z")], "Ok"
+        )
+        ref_backbone = reference_merge([dict(m) for m in pd], [dict(m) for m in pc], [dict(m) for m in pc] + [_msg("assistant", "z")], "Ok")
+        # Compare the backfilled visible backbone (role+content sequence).
+        opt_seq = [(m.get("role"), m.get("content")) for m in opt]
+        ref_seq = [(m.get("role"), m.get("content")) for m in ref_backbone]
+        # opt includes the appended tail delta; ref_backbone is backbone only — so
+        # ref must be a prefix-compatible subsequence. Assert backbone equality up to
+        # ref length.
+        assert opt_seq[: len(ref_seq)] == ref_seq, (
+            f"optimized backbone diverged from reference\n  pd={pd}\n  pc={pc}\n  opt={opt_seq}\n  ref={ref_seq}"
+        )