From e635bf370be43a21897aef140f658061bf15f357 Mon Sep 17 00:00:00 2001 From: Kris Urbas <605420+krzysu@users.noreply.github.com> Date: Mon, 20 Jul 2026 08:08:35 +0200 Subject: [PATCH 1/2] fix(gateway): suppress stream consumer on API failure and cap fallback chunks to prevent Telegram flood --- agent/agent_init.py | 7 + agent/conversation_loop.py | 16 ++ gateway/run.py | 12 ++ gateway/stream_consumer.py | 131 ++++++++++++++++ .../test_stream_consumer_flood_guard.py | 147 ++++++++++++++++++ 5 files changed, 313 insertions(+) create mode 100644 tests/gateway/test_stream_consumer_flood_guard.py diff --git a/agent/agent_init.py b/agent/agent_init.py index 407f9a6c7b4b6..23bcf740cc5ce 100644 --- a/agent/agent_init.py +++ b/agent/agent_init.py @@ -866,6 +866,13 @@ def init_agent( agent._last_activity_desc: str = "initializing" agent._current_tool: str | None = None agent._api_call_count: int = 0 + # Short summary of the model-API failure, set when the conversation loop + # exhausts retries on a terminal error (rate-limit / 429 / connection + # drop). The gateway's stream consumer reads this via ``api_failed_summary`` + # to suppress a partial / oversized streamed buffer and deliver a single + # clean error instead of flooding Telegram with split messages. ``None`` + # when the last turn succeeded (or hasn't failed yet). + agent.api_failed_summary: Optional[str] = None # Opt-out flag for the between-turns MCP tool refresh (build_turn_context). # Set on internal forks (e.g. background_review) that must keep ``tools[]`` # byte-identical to a parent for provider cache parity. diff --git a/agent/conversation_loop.py b/agent/conversation_loop.py index 3ca96898cf997..65e01fad9790d 100644 --- a/agent/conversation_loop.py +++ b/agent/conversation_loop.py @@ -1211,6 +1211,10 @@ def run_conversation( _last_preflight_pressure: Optional[int] = None _preflight_compression_blocked = _ctx.preflight_compression_blocked _turn_exit_reason = "unknown" # Diagnostic: why the loop ended + # Reset the model-API failure summary at the start of every turn so a + # stale error from a previous failed turn can never suppress a valid + # response in the gateway's stream consumer. + agent.api_failed_summary = None # Last composed answer intentionally held back by a verification gate. If # that continuation consumes the remaining budget, this is the best # user-facing result available; it must not be confused with error or @@ -5111,6 +5115,11 @@ def _perform_api_call(next_api_kwargs): # Terminal — flush buffered retry/fallback trace. agent._flush_status_buffer() _final_summary = agent._summarize_api_error(api_error) + # Surface the failure to the gateway's stream consumer so + # it can suppress a partial / oversized streamed buffer + # (e.g. echoed system prompt) and deliver one clean error + # instead of flooding the user with split messages. + agent.api_failed_summary = _final_summary _billing_guidance = "" if classified.reason == FailoverReason.billing: agent._emit_status(f"❌ Billing or credits exhausted — {_final_summary}") @@ -5442,6 +5451,13 @@ def _perform_api_call(next_api_kwargs): if response is None: _turn_exit_reason = "all_retries_exhausted_no_response" print(f"{agent.log_prefix}❌ All API retries exhausted with no successful response.") + # Surface the failure to the gateway's stream consumer so it can + # suppress a partial / oversized streamed buffer and deliver one + # clean error instead of flooding the user with split messages. + agent.api_failed_summary = ( + getattr(agent, "api_failed_summary", None) + or "All API retries exhausted with no successful response." + ) agent._persist_session(messages, conversation_history) break diff --git a/gateway/run.py b/gateway/run.py index 06a26d73aeabf..283b7bd585d1e 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -4191,6 +4191,18 @@ def run_sync(self): on_before_finalize=_pause_typing_before_finalize, initial_reply_to_id=ctx.event_message_id, run_still_current=ctx._run_still_current, + # Lazily report the agent's model-API failure so + # the stream consumer can suppress a partial / + # oversized streamed buffer (e.g. echoed system + # prompt) and deliver a single clean error instead + # of flooding the user with split messages. + api_error_fn=( + lambda: getattr( + ctx.agent_holder[0], "api_failed_summary", None + ) + if ctx.agent_holder and ctx.agent_holder[0] is not None + else None + ), ) if _want_stream_deltas: def _stream_delta_cb(text: str) -> None: diff --git a/gateway/stream_consumer.py b/gateway/stream_consumer.py index 0e601be801606..33c395c969b0c 100644 --- a/gateway/stream_consumer.py +++ b/gateway/stream_consumer.py @@ -172,6 +172,16 @@ class GatewayStreamConsumer: # progressive edits for the remainder of the stream. _MAX_FLOOD_STRIKES = 3 + # Hard ceiling on how many separate platform messages a fallback/overflow + # send may produce. When a response would split into more chunks than + # this, the consumer aborts the bulk send and instead delivers ONE short + # safe message. Without this cap a large accumulated buffer (e.g. the + # full system prompt echoed back, or an oversized model reply) becomes a + # flood of ~100 separate Telegram messages. 10 is well above any + # legitimate single reply that needs split delivery yet far below the + # flood territory. + _MAX_FALLBACK_CHUNKS = 10 + # Reasoning/thinking tags that models emit inline in content. # Must stay in sync with cli.py _OPEN_TAGS/_CLOSE_TAGS and # run_agent.py _strip_think_blocks() tag variants. @@ -199,6 +209,7 @@ def __init__( on_before_finalize: Optional[Callable[[], Any]] = None, initial_reply_to_id: Optional[str] = None, run_still_current: Optional[Callable[[], bool]] = None, + api_error_fn: Optional[Callable[[], Optional[str]]] = None, ): self.adapter = adapter self.chat_id = chat_id @@ -283,6 +294,17 @@ def __init__( # continuing to edit and deliver stale deltas. self._run_still_current = run_still_current or (lambda: True) + # API-failure detector. When the agent's model call fails + # (rate-limit / 429 / connection drop), this callable returns the + # short error summary (or a non-empty truthy string). The consumer + # then suppresses the accumulated streamed content at final flush and + # delivers just the clean error instead of blasting the partial + # buffer (which may contain the full system prompt / skill context) + # to the user — that buffer is exactly what produced the Telegram + # flood-of-messages symptom when the model API was down. Optional; + # consumers that don't wire it keep the legacy (dangerous) behaviour. + self._api_error_fn = api_error_fn + # Think-block filter state (mirrors CLI's _stream_delta tag suppression) self._in_think_block = False self._think_buffer = "" @@ -803,12 +825,35 @@ async def run(self) -> None: ): should_edit = False if should_edit and self._accumulated: + # API-failure guard (see _send_api_error_final): if the + # model call failed, never flush the accumulated buffer to + # the user — deliver only the clean error. Intercepted + # here so it covers both mid-stream overflow splits and + # the final flush at got_done. + _api_err = ( + self._api_error_fn() if self._api_error_fn else None + ) + if _api_err: + await self._send_api_error_final(_api_err) + return # Split overflow: if accumulated text exceeds the platform # limit, split into properly sized chunks. if ( _len_fn(self._accumulated) > _safe_limit and self._message_id is None ): + # Flood guard for fallback mode: when edits are broken + # (flood-control strikes exhausted / fallback promoted) + # a huge buffer would otherwise be chunked into dozens + # of separate messages. Refuse and deliver ONE safe + # message instead. + if self._fallback_final_send: + _cap_chunks = self.adapter.truncate_message( + self._accumulated, _safe_limit, len_fn=_len_fn, + ) + if len(_cap_chunks) > self._MAX_FALLBACK_CHUNKS: + await self._send_fallback_too_large() + return # No existing message to edit (first message or after a # segment break). Seal only the overflowing head chunks # as fixed messages, then keep the trailing chunk in @@ -949,6 +994,22 @@ async def run(self) -> None: if got_done: if self._accumulated or self._message_id is not None or self._already_sent: await self._notify_before_finalize() + + # API-failure guard: if the model call failed (rate-limit / + # 429 / connection drop), the accumulated buffer is partial, + # possibly huge (echoed system prompt / skill context), and + # must NOT be flushed to the user. Deliver only a single + # clean error message and mark the response as sent so the + # gateway's own final-send path skips re-delivering the raw + # buffer. This is the primary fix for the Telegram flood + # symptom when the model API is down. + _api_err = ( + self._api_error_fn() if self._api_error_fn else None + ) + if _api_err: + await self._send_api_error_final(_api_err) + return + # Final edit without cursor. If progressive editing failed # mid-stream, send a single continuation/fallback message # here instead of letting the base gateway path send the @@ -1242,6 +1303,66 @@ def _truncate_for_stream( return self._split_text_chunks(text, limit, len_fn) return list(chunks) + async def _send_api_error_final(self, error_summary: str) -> None: + """Deliver a single clean error message when the model API failed. + + Called from the final-flush path when ``api_error_fn`` reports the + agent's model call failed. The accumulated streamed buffer is + discarded (it may be partial and huge — echoed system prompt / skill + context) so we never flush it to the user. We mark the response as + delivered so the gateway's own final-send path skips re-delivering the + raw buffer (which would otherwise be chunked into another flood). + """ + # Keep it short and user-friendly. The long error detail stays in the + # gateway/agent logs; the user only needs to know to retry. + _safe = f"⚠️ Model API error — no response was generated. {error_summary}".strip() + try: + result = await self.adapter.send( + chat_id=self.chat_id, + content=_safe, + metadata=self._metadata_for_send(final=True), + ) + except Exception: + result = None + # Mark delivered regardless of send success: if the send itself failed + # (e.g. Telegram also unreachable), the gateway's fallback send will + # still try the agent's short error message. We never want the large + # accumulated buffer to be re-delivered. + self._already_sent = True + self._final_response_sent = True + self._final_content_delivered = True + if result and result.success and result.message_id: + self._message_id = str(result.message_id) + self._last_sent_text = _safe + + async def _send_fallback_too_large(self) -> None: + """Deliver a single safe message when a response is too large to send. + + Used by the fallback-send flood guard (``_MAX_FALLBACK_CHUNKS``). The + oversized accumulated buffer is discarded and replaced with one short + message so the user is never flooded with dozens of split messages. + Marked delivered so the gateway's normal final-send path does not + re-deliver the raw (equally oversized) buffer. + """ + _safe = ( + "⚠️ The response was too large to deliver safely. " + "Please try again — the full answer will be regenerated." + ) + try: + result = await self.adapter.send( + chat_id=self.chat_id, + content=_safe, + metadata=self._metadata_for_send(final=True), + ) + except Exception: + result = None + self._already_sent = True + self._final_response_sent = True + self._final_content_delivered = True + if result and result.success and result.message_id: + self._message_id = str(result.message_id) + self._last_sent_text = _safe + async def _send_fallback_final(self, text: str) -> None: """Send the final continuation after streaming edits stop working. @@ -1340,6 +1461,16 @@ async def _send_fallback_final(self, text: str) -> None: safe_limit = max(500, raw_limit - 100) chunks = self._split_text_chunks(continuation, safe_limit, len_fn=_len_fn) + # Flood guard: refuse to blast an oversized buffer as dozens of + # separate messages. When the response would split into more chunks + # than the cap, deliver ONE short safe message instead. This catches + # the non-API-error case (e.g. the model returns a huge reply and + # Telegram network flapping promotes us to fallback mode) so the user + # never gets a 100-message dump. + if len(chunks) > self._MAX_FALLBACK_CHUNKS: + await self._send_fallback_too_large() + return + stale_message_id = self._message_id # partial message to clean up last_message_id: Optional[str] = None last_successful_chunk = "" diff --git a/tests/gateway/test_stream_consumer_flood_guard.py b/tests/gateway/test_stream_consumer_flood_guard.py new file mode 100644 index 0000000000000..af54ab4fd4951 --- /dev/null +++ b/tests/gateway/test_stream_consumer_flood_guard.py @@ -0,0 +1,147 @@ +"""Tests for the stream-consumer flood guards. + +When the model API fails (rate-limit / 429 / connection drop) the +``GatewayStreamConsumer`` may hold a partial, possibly huge, accumulated +buffer (e.g. an echoed system prompt / skill context). Historically that +buffer was flushed to the user as a flood of split Telegram messages. + +These tests pin the two guards that prevent that: + +* ``api_error_fn`` — when the agent's model call failed, the consumer + suppresses the accumulated buffer and delivers ONE short clean error. +* ``_MAX_FALLBACK_CHUNKS`` — when a response would split into more chunks + than the cap, the consumer delivers ONE short safe message instead of a + flood. +""" + +from __future__ import annotations + +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock + +import pytest + +from gateway.stream_consumer import GatewayStreamConsumer, StreamConsumerConfig + + +def _make_adapter(*, supports_delete: bool = True) -> MagicMock: + """Minimal MagicMock adapter wired for send/edit/delete.""" + adapter = MagicMock() + adapter.REQUIRES_EDIT_FINALIZE = False + adapter.MAX_MESSAGE_LENGTH = 4096 + adapter.send = AsyncMock(return_value=SimpleNamespace( + success=True, message_id="preview_1", + )) + adapter.edit_message = AsyncMock(return_value=SimpleNamespace( + success=True, message_id="preview_1", + )) + if supports_delete: + adapter.delete_message = AsyncMock(return_value=True) + else: + del adapter.delete_message # type: ignore[attr-defined] + return adapter + + +def _sent_texts(adapter) -> list[str]: + texts = [] + for call in adapter.send.call_args_list: + texts.append(call.kwargs.get("content", "")) + if getattr(adapter, "edit_message", None) is not None: + for call in adapter.edit_message.call_args_list: + texts.append(call.kwargs.get("content", "")) + return texts + + +class TestApiErrorSuppression: + @pytest.mark.asyncio + async def test_api_error_suppresses_accumulated_buffer(self): + """On API failure the consumer sends the clean error, not the buffer. + + The buffer is held until ``got_done`` (large buffer_threshold) to + mimic the real fallback scenario: the consumer accumulates the whole + streamed response and would otherwise dump it as a flood of split + messages at finalization. + """ + adapter = _make_adapter() + # A big buffer that looks like an echoed system prompt + skills. + big_buffer = "SYSTEM PROMPT ... " + "skill content " * 5000 + consumer = GatewayStreamConsumer( + adapter, "chat_1", + StreamConsumerConfig(edit_interval=0.01, buffer_threshold=10_000_000), + api_error_fn=lambda: "HTTP 429: Weekly usage limit reached.", + ) + consumer.on_delta(big_buffer) + consumer.finish() + await consumer.run() + + sent = _sent_texts(adapter) + # Exactly one message, and it is the short clean error — never the + # raw buffer. + assert len(sent) == 1 + assert "HTTP 429" in sent[0] + assert big_buffer[:20] not in sent[0] + + # Delivery flags set so the gateway skips re-delivering the buffer. + assert consumer.final_response_sent is True + assert consumer.final_content_delivered is True + assert consumer.already_sent is True + + @pytest.mark.asyncio + async def test_no_api_error_keeps_legacy_behaviour(self): + """Without an api_error_fn the consumer still delivers the buffer.""" + adapter = _make_adapter() + consumer = GatewayStreamConsumer( + adapter, "chat_1", + StreamConsumerConfig(edit_interval=0.01, buffer_threshold=10_000_000), + ) + consumer.on_delta("Hello from the model") + consumer.finish() + await consumer.run() + + sent = _sent_texts(adapter) + assert any("Hello from the model" in t for t in sent) + # Delivery flags set (legacy behaviour) — no API-error short-circuit. + assert consumer.final_response_sent is True + + +class TestFallbackChunkCap: + @pytest.mark.asyncio + async def test_oversized_fallback_sends_single_safe_message(self): + """A buffer that would split past the cap becomes one safe message. + + We force fallback mode (edits unsupported) and feed a buffer far + larger than ``_MAX_FALLBACK_CHUNKS * MAX_MESSAGE_LENGTH``. + """ + from gateway.platforms.base import BasePlatformAdapter, SendResult + + # Use a real BasePlatformAdapter subclass so truncate_message behaves + # (a plain MagicMock's truncate_message returns a non-iterable stub). + CapAdapter = type("CapAdapter", (BasePlatformAdapter,), {"MAX_MESSAGE_LENGTH": 4096}) + CapAdapter.__abstractmethods__ = frozenset() + adapter = CapAdapter.__new__(CapAdapter) + adapter._typing_paused = set() + adapter.send = AsyncMock(return_value=SimpleNamespace(success=True, message_id="m1")) + # Edits fail → promotes the consumer into fallback mode. + adapter.edit_message = AsyncMock(return_value=SimpleNamespace(success=False, message_id=None)) + adapter.REQUIRES_EDIT_FINALIZE = False + + # Huge buffer that will split into far more than the cap of chunks. + huge = "x" * (GatewayStreamConsumer._MAX_FALLBACK_CHUNKS * 4096 * 3) + consumer = GatewayStreamConsumer( + adapter, "chat_1", + StreamConsumerConfig(edit_interval=0.01, buffer_threshold=10_000_000), + ) + # Enter fallback mode by exhausting flood strikes up front. + consumer._fallback_final_send = True + consumer.on_delta(huge) + consumer.finish() + await consumer.run() + + sent = _sent_texts(adapter) + # The flood guard must NOT have delivered the raw buffer as many + # messages — it sends exactly one safe message instead. + assert len(sent) == 1 + assert "too large" in sent[0] + assert huge[:20] not in sent[0] + assert consumer.final_response_sent is True + assert consumer.already_sent is True From fe41fa56e7b890332ab51d195f720c591c0adc07 Mon Sep 17 00:00:00 2001 From: Kris Urbas <605420+krzysu@users.noreply.github.com> Date: Thu, 30 Jul 2026 14:03:10 +0200 Subject: [PATCH 2/2] fix(gateway): drop unconditional fallback chunk cap per review --- gateway/stream_consumer.py | 60 ------ .../test_stream_consumer_flood_guard.py | 197 ++++++++++++++---- 2 files changed, 158 insertions(+), 99 deletions(-) diff --git a/gateway/stream_consumer.py b/gateway/stream_consumer.py index 33c395c969b0c..6de3f892fad19 100644 --- a/gateway/stream_consumer.py +++ b/gateway/stream_consumer.py @@ -172,16 +172,6 @@ class GatewayStreamConsumer: # progressive edits for the remainder of the stream. _MAX_FLOOD_STRIKES = 3 - # Hard ceiling on how many separate platform messages a fallback/overflow - # send may produce. When a response would split into more chunks than - # this, the consumer aborts the bulk send and instead delivers ONE short - # safe message. Without this cap a large accumulated buffer (e.g. the - # full system prompt echoed back, or an oversized model reply) becomes a - # flood of ~100 separate Telegram messages. 10 is well above any - # legitimate single reply that needs split delivery yet far below the - # flood territory. - _MAX_FALLBACK_CHUNKS = 10 - # Reasoning/thinking tags that models emit inline in content. # Must stay in sync with cli.py _OPEN_TAGS/_CLOSE_TAGS and # run_agent.py _strip_think_blocks() tag variants. @@ -842,18 +832,6 @@ async def run(self) -> None: _len_fn(self._accumulated) > _safe_limit and self._message_id is None ): - # Flood guard for fallback mode: when edits are broken - # (flood-control strikes exhausted / fallback promoted) - # a huge buffer would otherwise be chunked into dozens - # of separate messages. Refuse and deliver ONE safe - # message instead. - if self._fallback_final_send: - _cap_chunks = self.adapter.truncate_message( - self._accumulated, _safe_limit, len_fn=_len_fn, - ) - if len(_cap_chunks) > self._MAX_FALLBACK_CHUNKS: - await self._send_fallback_too_large() - return # No existing message to edit (first message or after a # segment break). Seal only the overflowing head chunks # as fixed messages, then keep the trailing chunk in @@ -1335,34 +1313,6 @@ async def _send_api_error_final(self, error_summary: str) -> None: self._message_id = str(result.message_id) self._last_sent_text = _safe - async def _send_fallback_too_large(self) -> None: - """Deliver a single safe message when a response is too large to send. - - Used by the fallback-send flood guard (``_MAX_FALLBACK_CHUNKS``). The - oversized accumulated buffer is discarded and replaced with one short - message so the user is never flooded with dozens of split messages. - Marked delivered so the gateway's normal final-send path does not - re-deliver the raw (equally oversized) buffer. - """ - _safe = ( - "⚠️ The response was too large to deliver safely. " - "Please try again — the full answer will be regenerated." - ) - try: - result = await self.adapter.send( - chat_id=self.chat_id, - content=_safe, - metadata=self._metadata_for_send(final=True), - ) - except Exception: - result = None - self._already_sent = True - self._final_response_sent = True - self._final_content_delivered = True - if result and result.success and result.message_id: - self._message_id = str(result.message_id) - self._last_sent_text = _safe - async def _send_fallback_final(self, text: str) -> None: """Send the final continuation after streaming edits stop working. @@ -1461,16 +1411,6 @@ async def _send_fallback_final(self, text: str) -> None: safe_limit = max(500, raw_limit - 100) chunks = self._split_text_chunks(continuation, safe_limit, len_fn=_len_fn) - # Flood guard: refuse to blast an oversized buffer as dozens of - # separate messages. When the response would split into more chunks - # than the cap, deliver ONE short safe message instead. This catches - # the non-API-error case (e.g. the model returns a huge reply and - # Telegram network flapping promotes us to fallback mode) so the user - # never gets a 100-message dump. - if len(chunks) > self._MAX_FALLBACK_CHUNKS: - await self._send_fallback_too_large() - return - stale_message_id = self._message_id # partial message to clean up last_message_id: Optional[str] = None last_successful_chunk = "" diff --git a/tests/gateway/test_stream_consumer_flood_guard.py b/tests/gateway/test_stream_consumer_flood_guard.py index af54ab4fd4951..42d40407cce8f 100644 --- a/tests/gateway/test_stream_consumer_flood_guard.py +++ b/tests/gateway/test_stream_consumer_flood_guard.py @@ -5,13 +5,18 @@ buffer (e.g. an echoed system prompt / skill context). Historically that buffer was flushed to the user as a flood of split Telegram messages. -These tests pin the two guards that prevent that: +These tests pin the guard that prevents that: * ``api_error_fn`` — when the agent's model call failed, the consumer suppresses the accumulated buffer and delivers ONE short clean error. -* ``_MAX_FALLBACK_CHUNKS`` — when a response would split into more chunks - than the cap, the consumer delivers ONE short safe message instead of a - flood. + +They also pin the **end-to-end wiring** the gateway relies on: the agent +exposes ``api_failed_summary`` after a terminal API failure, and the +lambda ``gateway/run.py`` uses to bridge that to the consumer returns +the summary at the moment the consumer checks it. This guards against +regressions where the agent-state lifecycle or the gateway wiring is +broken (a direct unit test on the consumer that injects a static +``api_error_fn`` would miss those). """ from __future__ import annotations @@ -24,7 +29,7 @@ from gateway.stream_consumer import GatewayStreamConsumer, StreamConsumerConfig -def _make_adapter(*, supports_delete: bool = True) -> MagicMock: +def _make_adapter() -> MagicMock: """Minimal MagicMock adapter wired for send/edit/delete.""" adapter = MagicMock() adapter.REQUIRES_EDIT_FINALIZE = False @@ -35,10 +40,7 @@ def _make_adapter(*, supports_delete: bool = True) -> MagicMock: adapter.edit_message = AsyncMock(return_value=SimpleNamespace( success=True, message_id="preview_1", )) - if supports_delete: - adapter.delete_message = AsyncMock(return_value=True) - else: - del adapter.delete_message # type: ignore[attr-defined] + adapter.delete_message = AsyncMock(return_value=True) return adapter @@ -46,12 +48,26 @@ def _sent_texts(adapter) -> list[str]: texts = [] for call in adapter.send.call_args_list: texts.append(call.kwargs.get("content", "")) - if getattr(adapter, "edit_message", None) is not None: - for call in adapter.edit_message.call_args_list: - texts.append(call.kwargs.get("content", "")) + for call in adapter.edit_message.call_args_list: + texts.append(call.kwargs.get("content", "")) return texts +def _gateway_api_error_lambda(agent_holder): + """Mirror the lambda wired by ``gateway/run.py`` into the consumer. + + ``gateway/run.py`` stores the agent in a one-element list and exposes + ``api_failed_summary`` via ``getattr(agent, "api_failed_summary", None)`` + so the consumer reads a *live* view of the agent state at the moment + of the final flush — not a snapshot taken at consumer construction. + """ + return lambda: ( + getattr(agent_holder[0], "api_failed_summary", None) + if agent_holder and agent_holder[0] is not None + else None + ) + + class TestApiErrorSuppression: @pytest.mark.asyncio async def test_api_error_suppresses_accumulated_buffer(self): @@ -104,44 +120,147 @@ async def test_no_api_error_keeps_legacy_behaviour(self): assert consumer.final_response_sent is True -class TestFallbackChunkCap: +class TestApiErrorEndToEndWiring: + """End-to-end coverage of the agent-state → consumer bridge. + + The agent records ``api_failed_summary`` on a terminal model failure, + and the gateway's ``api_error_fn`` lambda reads it lazily. This + exercises the full chain through the real lambda shape (not a static + callable) so a regression in either the agent-state lifecycle OR the + gateway wiring surfaces here. + """ + @pytest.mark.asyncio - async def test_oversized_fallback_sends_single_safe_message(self): - """A buffer that would split past the cap becomes one safe message. + async def test_terminal_api_failure_is_suppressed_at_final_flush(self): + """A large buffer accumulated BEFORE the agent reports failure is suppressed. - We force fallback mode (edits unsupported) and feed a buffer far - larger than ``_MAX_FALLBACK_CHUNKS * MAX_MESSAGE_LENGTH``. + Reproduces the real Telegram flood symptom: the model returns + partial content (often echoing the system prompt / skill context) + before the conversation loop records the terminal API error. When + ``got_done`` arrives, the consumer must consult the live + ``api_failed_summary`` and suppress the buffer. """ - from gateway.platforms.base import BasePlatformAdapter, SendResult - - # Use a real BasePlatformAdapter subclass so truncate_message behaves - # (a plain MagicMock's truncate_message returns a non-iterable stub). - CapAdapter = type("CapAdapter", (BasePlatformAdapter,), {"MAX_MESSAGE_LENGTH": 4096}) - CapAdapter.__abstractmethods__ = frozenset() - adapter = CapAdapter.__new__(CapAdapter) - adapter._typing_paused = set() - adapter.send = AsyncMock(return_value=SimpleNamespace(success=True, message_id="m1")) - # Edits fail → promotes the consumer into fallback mode. - adapter.edit_message = AsyncMock(return_value=SimpleNamespace(success=False, message_id=None)) - adapter.REQUIRES_EDIT_FINALIZE = False - - # Huge buffer that will split into far more than the cap of chunks. - huge = "x" * (GatewayStreamConsumer._MAX_FALLBACK_CHUNKS * 4096 * 3) + agent_holder: list = [None] + + class _StubAgent: + # The agent's __init__ signature is irrelevant — the gateway + # only ever reads ``api_failed_summary`` off the live object. + api_failed_summary = None + + agent_holder[0] = _StubAgent() + + adapter = _make_adapter() + big_buffer = "SYSTEM PROMPT ... " + "skill content " * 5000 consumer = GatewayStreamConsumer( adapter, "chat_1", StreamConsumerConfig(edit_interval=0.01, buffer_threshold=10_000_000), + api_error_fn=_gateway_api_error_lambda(agent_holder), ) - # Enter fallback mode by exhausting flood strikes up front. - consumer._fallback_final_send = True - consumer.on_delta(huge) + # Stream partial content first (mimics a model that returns some + # output then errors mid-stream). + consumer.on_delta(big_buffer) + # The conversation loop's terminal error path records the summary + # *between* the last delta and the consumer's final flush — exactly + # when ``gateway/run.py`` is calling the lambda. + agent_holder[0].api_failed_summary = "HTTP 429: Weekly usage limit reached." consumer.finish() await consumer.run() sent = _sent_texts(adapter) - # The flood guard must NOT have delivered the raw buffer as many - # messages — it sends exactly one safe message instead. + # Exactly one message, the clean error — never the raw buffer. assert len(sent) == 1 - assert "too large" in sent[0] - assert huge[:20] not in sent[0] + assert "HTTP 429" in sent[0] + assert big_buffer[:20] not in sent[0] assert consumer.final_response_sent is True + assert consumer.final_content_delivered is True assert consumer.already_sent is True + + @pytest.mark.asyncio + async def test_no_failure_does_not_suppress_streamed_content(self): + """When the agent never reports failure, the guard does not fire. + + Guards against an over-eager guard that would suppress legitimate + large replies: ``_send_api_error_final`` must not run, so no + ``"⚠️ Model API error"`` text reaches the user and the consumer + does not pre-emptively set ``_final_content_delivered`` (which + would otherwise tell the gateway to skip its own final send). + """ + agent_holder: list = [None] + + class _StubAgent: + api_failed_summary = None + + agent_holder[0] = _StubAgent() + + adapter = _make_adapter() + consumer = GatewayStreamConsumer( + adapter, "chat_1", + StreamConsumerConfig(edit_interval=0.01, buffer_threshold=10_000_000), + api_error_fn=_gateway_api_error_lambda(agent_holder), + ) + consumer.on_delta("Here is a complete, well-formed reply.") + consumer.finish() + await consumer.run() + + sent = _sent_texts(adapter) + # The clean error must never appear — only the legitimate content. + assert not any("Model API error" in t for t in sent) + # And the actual reply must reach the user. + assert any("complete, well-formed reply" in t for t in sent) + + @pytest.mark.asyncio + async def test_summary_set_after_buffer_but_before_flush_is_seen(self): + """A late ``api_failed_summary`` (set after the last delta) is honored. + + In the real flow, the conversation loop records the terminal error + *during* the API call and the consumer hasn't run its final flush + yet. A summary that appears between the last ``on_delta`` and + ``run()`` must still suppress the buffer. + """ + agent_holder: list = [None] + + class _StubAgent: + api_failed_summary = None + + agent_holder[0] = _StubAgent() + + adapter = _make_adapter() + consumer = GatewayStreamConsumer( + adapter, "chat_1", + StreamConsumerConfig(edit_interval=0.01, buffer_threshold=10_000_000), + api_error_fn=_gateway_api_error_lambda(agent_holder), + ) + # Buffer accumulated cleanly; only later does the agent learn it failed. + consumer.on_delta("looks like a normal response ... " * 500) + agent_holder[0].api_failed_summary = ( + "All API retries exhausted with no successful response." + ) + consumer.finish() + await consumer.run() + + sent = _sent_texts(adapter) + assert len(sent) == 1 + assert "retries exhausted" in sent[0] + assert "looks like a normal response" not in sent[0] + + @pytest.mark.asyncio + async def test_lambda_handles_missing_agent_holder(self): + """If the agent slot is empty, the lambda returns None (no suppression). + + Defends against a regression where the gateway wires the lambda + before the agent is created. The consumer must fall back to its + legacy (non-suppressing) behaviour. + """ + agent_holder: list = [None] + adapter = _make_adapter() + consumer = GatewayStreamConsumer( + adapter, "chat_1", + StreamConsumerConfig(edit_interval=0.01, buffer_threshold=10_000_000), + api_error_fn=_gateway_api_error_lambda(agent_holder), + ) + consumer.on_delta("normal response") + consumer.finish() + await consumer.run() + + sent = _sent_texts(adapter) + assert any("normal response" in t for t in sent)