Cap a message by a shared deadline and bound streaming's first chunk by it

Lock wait, search, route retry and history summary each had their own timeout, so one message could hold the chat lock for minutes. stream_interrupted now survives a successful return so continue works.

bot.py

@@ -86,6 +86,23 @@ def _current_log_secrets() -> tuple[str, ...]:
 # id, и собрать путь одного запроса по логу было нельзя (прод 27.10.2026).
 _LOG_CHAT: contextvars.ContextVar[str] = contextvars.ContextVar("log_chat", default="")
 
+# Общий дедлайн сообщения: каждый этап берёт min(свой таймаут, остаток), иначе
+# лок + поиск + маршрут + саммари в сумме давали до двух минут (аудит 27.10.2026).
+_MSG_DEADLINE: contextvars.ContextVar[float | None] = contextvars.ContextVar("msg_deadline", default=None)
+
+
+def _msg_remaining_sec() -> float | None:
+    """Сколько осталось до потолка сообщения; None — обработка вне handle_message."""
+    deadline = _MSG_DEADLINE.get()
+    return None if deadline is None else max(0.0, deadline - time.monotonic())
+
+
+def _deadline_within_message(seconds: float) -> float:
+    """Дедлайн этапа, не выходящий за общий потолок сообщения."""
+    now = time.monotonic()
+    remaining = _msg_remaining_sec()
+    return now + seconds if remaining is None else now + min(seconds, remaining)
+
 
 class _SecretLogFormatter(logging.Formatter):
     def format(self, record: logging.LogRecord, _secrets=_current_log_secrets) -> str:
@@ -405,6 +422,9 @@ ROUTE_MODEL_TIMEOUT_SEC = _env_number("ROUTE_MODEL_TIMEOUT_SEC", 22, min_value=1
 # ROUTE_TOTAL_BUDGET_SEC — общий бюджет на ВЕСЬ маршрут одного сообщения. Без него
 # каскадный сбой у многих моделей растягивал бы один ответ на минуты под локом чата.
 ROUTE_TOTAL_BUDGET_SEC = _env_number("ROUTE_TOTAL_BUDGET_SEC", 40, min_value=1)
+# MESSAGE_TOTAL_BUDGET_SEC — потолок на всё сообщение (лок + поиск + маршрут).
+# Этапы ниже берут min(свой таймаут, остаток), иначе их сумма выходила за бюджет.
+MESSAGE_TOTAL_BUDGET_SEC = _env_number("MESSAGE_TOTAL_BUDGET_SEC", 90, min_value=10)
 # HISTORY_SUMMARY_BUDGET_SEC — потолок саммаризации истории: без него висящие
 # вызовы держали бы лок чата десятки секунд уже после готового ответа.
 HISTORY_SUMMARY_BUDGET_SEC = _env_number("HISTORY_SUMMARY_BUDGET_SEC", 30, min_value=1)
@@ -1301,10 +1321,13 @@ async def handle_message(message: Message) -> None:
     # Метка чата на всех строках обработки (см. _LOG_CHAT); сброс в finally —
     # иначе значение утечёт в следующий апдейт этого же процесса.
     token = _LOG_CHAT.set(f"chat={message.chat.id if message.chat else 0}")
+    # Потолок сообщения ставится здесь: альбомные задачи наследуют контекст.
+    dtoken = _MSG_DEADLINE.set(time.monotonic() + MESSAGE_TOTAL_BUDGET_SEC)
     try:
         await _handle_message(message)
     finally:
         _LOG_CHAT.reset(token)
+        _MSG_DEADLINE.reset(dtoken)
 
 
 async def _handle_message(message: Message) -> None:
@@ -1339,8 +1362,10 @@ async def _handle_message(message: Message) -> None:
         await _handle_message_core(message)
         return
 
+    remaining = _msg_remaining_sec()
+    lock_timeout = CHAT_LOCK_TIMEOUT_SEC if remaining is None else max(1.0, min(CHAT_LOCK_TIMEOUT_SEC, remaining))
     try:
-        lock = await acquire_chat_lock(chat_id, CHAT_LOCK_TIMEOUT_SEC)
+        lock = await acquire_chat_lock(chat_id, lock_timeout)
     except asyncio.TimeoutError:
         log.warning("[lock] Timeout waiting for lock on chat %s", chat_id)
         with contextlib.suppress(Exception):

lumen_chat_state.py

@@ -1325,9 +1325,14 @@ async def _trim_history(history: list[dict[str, Any]]) -> None:
     digest_input = "\n".join(lines)[:_HISTORY_SUMMARY_INPUT_MAX_CHARS]
     summary = ""
     if digest_input.strip():
+        # Саммари не должно выходить за потолок сообщения: оно идёт под локом чата.
+        budget = bot.HISTORY_SUMMARY_BUDGET_SEC
+        remaining = bot._msg_remaining_sec()
+        if remaining is not None:
+            budget = max(0.1, min(budget, remaining))
         try:
             summary = await asyncio.wait_for(
-                _summarize_text(digest_input), timeout=bot.HISTORY_SUMMARY_BUDGET_SEC,
+                _summarize_text(digest_input), timeout=budget,
             )
         except Exception as exc:
             log.warning("[history] Summarization over budget/failed, plain cut: %s", exc)

lumen_streaming.py

@@ -251,9 +251,14 @@ async def _run_streaming_reply(
         bot._record_stats_event("model_attempts")
 
         # Дальше цикл читает генератор напрямую — проверки кусков/утечек/троттлинга ниже без изменений.
+        first_chunk_limit = _model_first_chunk_limit(_model_speed_key(provider, model_id), bot.FIRST_CHUNK_TIMEOUT_SEC)
+        if deadline is not None:
+            # Первый кусок не должен ждать дольше общего бюджета сообщения: до него
+            # дедлайн не проверялся и ожидание могло выйти за бюджет (аудит 27.10.2026).
+            first_chunk_limit = min(first_chunk_limit, max(0.5, deadline - time.monotonic()))
         piece_agen = _pieces_with_waiting_feedback(
             piece_agen, placeholder,
-            first_chunk_limit=_model_first_chunk_limit(_model_speed_key(provider, model_id), bot.FIRST_CHUNK_TIMEOUT_SEC),
+            first_chunk_limit=first_chunk_limit,
         )
         async for piece in piece_agen:
             if not piece:
@@ -433,7 +438,10 @@ async def _run_streaming_reply(
                 await bot._tg_call(sent_messages[-1].edit_text, _md_to_html(final_text + note), parse_mode=ParseMode.HTML, call_timeout=15.0)
             # Частичный ответ уже показан и ляжет в историю ниже — запоминаем, что
             # его оборвало: слово «продолжи» добьёт с места обрыва (_continue_after_interrupt).
+            # stream_interrupted переживает успешный возврат: иначе _handle_message_core
+            # снимал общий флаг и «продолжи» не срабатывал (аудит 27.10.2026).
             state["interrupted"] = True
+            state["stream_interrupted"] = True
             bot.mark_state_dirty(chat_id)
     finally:
         aclose = getattr(piece_agen, "aclose", None)

tests/test_bot_streaming.py

@@ -581,6 +581,36 @@ def test_streaming_respects_route_deadline():
         bot.chat_state.pop(chat_id, None)
 
 
+def test_streaming_first_chunk_limit_capped_by_route_deadline(monkeypatch):
+    # Что защищает: ожидание первого куска не выходит за бюджет маршрута.
+    # Регрессия: предел модели (по замерам до 75с) игнорировал дедлайн, и стрим
+    # держал лок сверх бюджета (аудит 27.10.2026).
+    import lumen_streaming
+    seen = {}
+
+    async def fake_feedback(piece_agen, placeholder, *, first_chunk_limit):
+        seen["limit"] = first_chunk_limit
+        async for piece in piece_agen:
+            yield piece
+
+    async def one_piece():
+        yield "готово"
+
+    monkeypatch.setattr(lumen_streaming, "_pieces_with_waiting_feedback", fake_feedback)
+    monkeypatch.setattr(lumen_streaming, "_model_first_chunk_limit", lambda key, floor: 3600)
+    chat_id = 999318
+    incoming = _FakeIncomingMessage(chat_id)
+    try:
+        answer, _ = asyncio.run(bot._run_streaming_reply(
+            chat_id, "Привет!", incoming, provider="openrouter", model_id="w:free",
+            piece_agen=one_piece(), deadline=time.monotonic() + 1,
+        ))
+        assert answer == "готово"
+        assert seen["limit"] <= 1.1
+    finally:
+        bot.chat_state.pop(chat_id, None)
+
+
 def _streaming_rate_limit_quotas(chat_id, model_id, error_text):
     async def pieces_429():
         raise bot.OpenRouterAPIError(error_text, status_code=429)

tests/test_bot_transport.py

@@ -525,3 +525,34 @@ def test_telegram_api_call_ok_false_reports_description(monkeypatch):
     assert "400" in msg
     assert "'ok'" not in msg
 
+
+def test_message_deadline_caps_stage_budget():
+    # Что защищает: общий потолок сообщения (MESSAGE_TOTAL_BUDGET_SEC). Регрессия:
+    # таймауты этапов складывались без учёта остатка, и обработка выходила за
+    # бюджет (аудит 27.10.2026).
+    token = bot._MSG_DEADLINE.set(time.monotonic() + 5)
+    try:
+        assert bot._msg_remaining_sec() <= 5
+        assert bot._deadline_within_message(40) <= time.monotonic() + 5.1
+    finally:
+        bot._MSG_DEADLINE.reset(token)
+    assert bot._msg_remaining_sec() is None
+    assert bot._deadline_within_message(40) > time.monotonic() + 39
+
+
+def test_handle_message_sets_and_resets_message_deadline(monkeypatch):
+    # Что защищает: установку потолка на входе в обработчик. Регрессия: дедлайн
+    # оставался бы в контексте после обработки и урезал следующий апдейт
+    # (аудит 27.10.2026).
+    from types import SimpleNamespace
+    seen = {}
+
+    async def fake_handle(message):
+        seen["remaining"] = bot._msg_remaining_sec()
+
+    monkeypatch.setattr(bot, "_handle_message", fake_handle)
+    asyncio.run(bot.handle_message(SimpleNamespace(chat=None, from_user=None)))
+    assert seen["remaining"] is not None and seen["remaining"] > 0
+    assert bot._msg_remaining_sec() is None
+
+