diff --git a/CHANGELOG.md b/CHANGELOG.md index 2d195a7d..36cfdd49 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/api/streaming.py b/api/streaming.py index 44c1b548..89e1417c 100644 --- a/api/streaming.py +++ b/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. diff --git a/tests/test_merge_backfill_perf_optimization.py b/tests/test_merge_backfill_perf_optimization.py new file mode 100644 index 00000000..4a21e424 --- /dev/null +++ b/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}" + )