fix(gateway): stop truncate_message hanging on a pathologically small max_length

BasePlatformAdapter.truncate_message() splits an over-length reply into
chunks. When max_length is 0 or 1 (and the content is longer), the split
loop makes no progress and spins forever, appending empty chunks — an
unbounded hang that pins a CPU and grows the chunk list until OOM:

  - headroom = max_length - INDICATOR_RESERVE - ... goes negative, and the
    < 1 fallback (max_length // 2) is also 0;
  - so _cp_limit is 0, the region is empty, no split point is found, and
    split_at falls back to _cp_limit (0);
  - chunk_body is remaining[:0] = "", remaining never shrinks, loop repeats.

The same stall is reachable under utf16_len (Telegram) whenever the next
char is a surrogate-pair emoji wider than the whole budget, so _cp_limit
maps to 0 codepoints even for max_length >= 2.

A pathological max_length is not hypothetical: the relay capability
descriptor's max_message_length is taken verbatim from the connector
(gateway/relay/descriptor.py from_json) and assigned straight to the
adapter's MAX_MESSAGE_LENGTH (gateway/relay/adapter.py), and 0 is a
documented "no limit" value there.

Guarantee forward progress: floor headroom at 1, and floor the
final split_at at max(1, _cp_limit) so at least one codepoint is always
consumed per iteration. Normal splitting is unaffected (both floors only
bite when the budget is already degenerate).

Adds regression tests that run truncate_message on a worker thread and
fail if it doesn't return: max_length 0/1/2 terminate and preserve every
character, and the utf16 emoji case terminates too.
This commit is contained in:
Drexuxux 2026-07-16 03:48:38 +03:00 committed by Teknium
parent 9fc8fe2176
commit fbf5005a7e
2 changed files with 59 additions and 2 deletions

View file

@ -5661,7 +5661,10 @@ class BasePlatformAdapter(ABC):
# a potential closing fence, and the chunk indicator.
headroom = max_length - INDICATOR_RESERVE - _len(prefix) - _len(FENCE_CLOSE)
if headroom < 1:
headroom = max_length // 2
# Floor at 1 so a pathologically small max_length (0 or 1 —
# e.g. a relay capability descriptor whose max_message_length
# is 0/1) can't make headroom 0 and stall the loop below.
headroom = max(1, max_length // 2)
# Everything remaining fits in one final chunk
if _len(prefix) + _len(remaining) <= max_length - INDICATOR_RESERVE:
@ -5685,7 +5688,13 @@ class BasePlatformAdapter(ABC):
if split_at < _cp_limit // 2:
split_at = region.rfind(" ")
if split_at < 1:
split_at = _cp_limit
# Consume at least one codepoint. Without the max(1, …) floor,
# a zero _cp_limit — reachable when max_length is 0/1, or under
# utf16_len when the next char is a surrogate pair wider than
# the whole budget — leaves split_at at 0, so ``remaining``
# never shrinks and the while-loop spins forever appending
# empty chunks (an unbounded hang / OOM).
split_at = max(1, _cp_limit)
# Avoid splitting inside an inline code span (`...`).
# If the text before split_at has an odd number of unescaped

View file

@ -1603,6 +1603,54 @@ class TestTruncateMessage:
assert "(1/" in chunks[0]
assert f"({len(chunks)}/{len(chunks)})" in chunks[-1]
@staticmethod
def _truncate_with_timeout(content, max_length, *, len_fn=None, timeout=3.0):
"""Run truncate_message on a worker thread; fail if it doesn't return.
Guards against the regression where a pathologically small max_length
made the split loop never consume any input and spin forever.
"""
import threading
box: dict = {}
def _run():
box["result"] = BasePlatformAdapter.truncate_message(
content, max_length, len_fn=len_fn
)
t = threading.Thread(target=_run, daemon=True)
t.start()
t.join(timeout=timeout)
assert not t.is_alive(), (
f"truncate_message hung (infinite loop) for max_length={max_length}"
)
return box["result"]
def test_pathological_small_max_length_terminates(self):
# max_length 0 and 1 previously drove the split loop into an unbounded
# hang (headroom -> 0, split_at -> 0, remaining never shrinks). It must
# terminate and preserve every character across the chunks.
import re
for max_length in (0, 1, 2):
chunks = self._truncate_with_timeout("abcdefghij", max_length)
assert chunks, f"no chunks for max_length={max_length}"
reassembled = "".join(
re.sub(r"\s*\(\d+/\d+\)$", "", c) for c in chunks
)
for ch in "abcdefghij":
assert ch in reassembled, f"char {ch!r} lost at max_length={max_length}"
def test_pathological_small_max_length_utf16_terminates(self):
# Under utf16_len (Telegram), a surrogate-pair emoji is 2 units wide, so
# a budget below that maps to zero codepoints — the same stall vector.
from gateway.platforms.base import utf16_len
chunks = self._truncate_with_timeout("😀😀😀😀😀", 1, len_fn=utf16_len)
assert chunks
assert "😀" in "".join(chunks)
def test_code_block_first_chunk_closed(self):
adapter = self._adapter()
msg = "Before\n```python\n" + "x = 1\n" * 100 + "```\nAfter"