diff --git a/gateway/run.py b/gateway/run.py index a5fe58f97a54..6a175cba0f23 100644 --- a/gateway/run.py +++ b/gateway/run.py @@ -16144,6 +16144,15 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew if event_type not in {"tool.started",}: return + # Suppress the redundant clarify progress bubble. The clarify poll + # itself (sent via send_clarify) already shows the question text and + # options, so a "❓ clarify: ..." status line above it is pure noise + # — and worse, it lands BETWEEN the flushed prose and the poll. The + # interactive prompt is the user-facing surface; the tool-progress + # echo adds nothing. + if tool_name == "clarify": + return + # Suppress tool-progress bubbles once the user has sent `stop`. # When the LLM response carries N parallel tool calls, the agent # fires N "tool.started" events back-to-back before checking for @@ -17244,6 +17253,25 @@ class GatewayRunner(GatewayAuthorizationMixin, GatewayKanbanWatchersMixin, Gatew except Exception: pass + # Ordering barrier (#clarify-ordering): flush any buffered + # assistant prose (interim commentary / streamed deltas) to the + # platform BEFORE sending the poll. The poll is delivered on a + # separate, agent-thread-blocking path; without this barrier it + # races ahead of prose still sitting in the stream consumer's + # queue, so the question renders ABOVE its own explanation. + # Best-effort + short timeout: never hang the agent thread if + # the consumer task isn't running. + try: + _sc = stream_consumer_holder[0] if stream_consumer_holder else None + _flush = getattr(_sc, "flush_pending_sync", None) + if callable(_flush): + _flush(timeout=3.0) + except Exception: + logger.debug( + "Stream-consumer flush before clarify prompt failed", + exc_info=True, + ) + send_ok = False fut = safe_schedule_threadsafe( _status_adapter.send_clarify( diff --git a/gateway/stream_consumer.py b/gateway/stream_consumer.py index a08e169f2f99..0eadcd931e72 100644 --- a/gateway/stream_consumer.py +++ b/gateway/stream_consumer.py @@ -20,6 +20,7 @@ import inspect import logging import queue import re +import threading import time from dataclasses import dataclass from typing import Any, Callable, Optional @@ -50,6 +51,15 @@ _NEW_SEGMENT = object() # API/tool iterations (for example: "I'll inspect the repo first."). _COMMENTARY = object() +# Queue marker for a synchronous flush barrier. Enqueued as +# ``(_FLUSH, threading.Event)``; the drain loop finalizes and delivers any +# buffered segment, then sets the event. A caller on the agent worker thread +# uses this (via ``flush_pending_sync``) to block until everything queued +# BEFORE the marker has actually landed on the platform — needed before +# sending a blocking interactive prompt (clarify poll) so the prompt is the +# last thing on screen, not racing ahead of buffered prose. +_FLUSH = object() + @dataclass class StreamConsumerConfig: @@ -321,6 +331,29 @@ class GatewayStreamConsumer: if text: self._queue.put((_COMMENTARY, text)) + def flush_pending_sync(self, timeout: float = 5.0) -> bool: + """Block the calling (agent worker) thread until everything queued + before this point has been finalized and delivered to the platform. + + Enqueues a ``(_FLUSH, Event)`` barrier behind any pending deltas / + commentary / segment breaks. The async ``run()`` task processes those + first (FIFO), then handles the barrier — finalizing the current segment + and setting the event. Returns True if the flush completed within + ``timeout``, False on timeout (so the caller continues rather than + hanging if the consumer task is not running / already finished). + + This is the ordering barrier used before sending a blocking interactive + prompt (clarify poll): without it, the poll — sent on a separate, + agent-thread-blocking path — races ahead of buffered prose that is still + sitting in this queue, so the question lands ABOVE its own explanation. + """ + evt = threading.Event() + try: + self._queue.put((_FLUSH, evt)) + except Exception: + return False + return evt.wait(timeout=max(0.0, float(timeout))) + def _notify_new_message(self) -> None: """Fire the on_new_message callback, swallowing any errors.""" cb = self._on_new_message @@ -331,6 +364,23 @@ class GatewayStreamConsumer: except Exception: logger.debug("on_new_message callback error", exc_info=True) + @staticmethod + def _signal_flush(flush_event) -> None: + """Wake a thread blocked in flush_pending_sync(), swallowing errors. + + Centralised so every loop exit path that consumed a ``_FLUSH`` barrier + (the normal bottom-of-iteration path AND early ``continue`` paths such + as the oversized-prose overflow split) reliably sets the event. Missing + a set is not a deadlock — the caller uses a bounded timeout — but it + would make the caller stall the full timeout before its blocking send. + """ + if flush_event is None: + return + try: + flush_event.set() + except Exception: + pass + def _reset_segment_state(self, *, preserve_no_edit: bool = False) -> None: if preserve_no_edit and self._message_id == "__no_edit__": return @@ -571,6 +621,8 @@ class GatewayStreamConsumer: # Drain all available items from the queue got_done = False got_segment_break = False + got_flush = False + flush_event = None commentary_text = None while True: try: @@ -584,6 +636,14 @@ class GatewayStreamConsumer: if isinstance(item, tuple) and len(item) == 2 and item[0] is _COMMENTARY: commentary_text = item[1] break + if isinstance(item, tuple) and len(item) == 2 and item[0] is _FLUSH: + # Flush barrier: finalize the current segment like a + # tool boundary, then signal the waiting thread once + # delivery for this iteration has completed (below). + got_flush = True + got_segment_break = True + flush_event = item[1] + break self._filter_and_accumulate(item) except queue.Empty: break @@ -689,8 +749,15 @@ class GatewayStreamConsumer: self._message_id = None self._fallback_final_send = False self._fallback_prefix = "" - continue + # This iteration consumed a _FLUSH barrier and delivered + # the buffered prose via the chunk loop above, then takes + # an early `continue` that skips the bottom-of-loop set. + # Signal here so flush_pending_sync() doesn't stall the + # full timeout waiting on already-delivered content. + if got_flush: + self._signal_flush(flush_event) + continue # Existing message: edit it with the first chunk, then # start a new message for the overflow remainder. while ( @@ -847,6 +914,13 @@ class GatewayStreamConsumer: await self._flush_segment_tail_on_edit_failure() self._reset_segment_state(preserve_no_edit=True) + # Flush barrier satisfied: the buffered segment (if any) has now + # been finalized and delivered above, so wake the thread blocked + # in flush_pending_sync(). Done last so the waiter only unblocks + # once everything queued before the barrier is on screen. + if got_flush: + self._signal_flush(flush_event) + await asyncio.sleep(0.05) # Small yield to not busy-loop except asyncio.CancelledError: @@ -879,6 +953,26 @@ class GatewayStreamConsumer: self._final_content_delivered = True except Exception as e: logger.error("Stream consumer error: %s", e) + finally: + # Safety net: if run() exits (normal return, cancellation, or + # exception) while a _FLUSH barrier is still queued or was consumed + # but not yet signaled, wake any waiters now. Without this a caller + # blocked in flush_pending_sync() would stall the full timeout when + # the consumer dies mid-flush. Bounded either way, but this makes + # the common case instant instead of timeout-delayed. + try: + while True: + item = self._queue.get_nowait() + if ( + isinstance(item, tuple) + and len(item) == 2 + and item[0] is _FLUSH + ): + self._signal_flush(item[1]) + except queue.Empty: + pass + except Exception: + pass # Strip MEDIA: tags before display. Uses the shared anchored # MEDIA_TAG_CLEANUP_RE from gateway/platforms/base.py — only tags whose diff --git a/tests/gateway/test_stream_consumer.py b/tests/gateway/test_stream_consumer.py index cd49d3d74782..7c8a783e79f8 100644 --- a/tests/gateway/test_stream_consumer.py +++ b/tests/gateway/test_stream_consumer.py @@ -2363,3 +2363,173 @@ class TestStripOrphanCloseTags: assert tag not in consumer._accumulated assert "trailing prose" in consumer._accumulated assert "more" in consumer._accumulated + + +# ── Flush barrier (clarify-ordering) tests ─────────────────────────────── + + +class TestFlushPendingSync: + """flush_pending_sync() is the ordering barrier that guarantees buffered + prose lands on the platform BEFORE a blocking interactive prompt (clarify + poll) is sent. Regression coverage for the bug where the poll raced ahead + of its own explanation, rendering the question above the prose. + """ + + @pytest.mark.asyncio + async def test_flush_delivers_buffered_commentary_before_returning(self): + """A commentary message queued before the flush barrier must be sent + to the adapter before flush_pending_sync() returns True.""" + adapter = MagicMock() + sent_order = [] + + async def _send(*args, **kwargs): + sent_order.append(("send", kwargs.get("content", ""))) + return SimpleNamespace(success=True, message_id="msg_1") + + async def _edit(*args, **kwargs): + sent_order.append(("edit", kwargs.get("content", ""))) + return SimpleNamespace(success=True) + + adapter.send = AsyncMock(side_effect=_send) + adapter.edit_message = AsyncMock(side_effect=_edit) + adapter.MAX_MESSAGE_LENGTH = 4096 + + config = StreamConsumerConfig(edit_interval=0.01, buffer_threshold=5) + consumer = GatewayStreamConsumer(adapter, "chat_123", config) + + # Mirror the live config that exhibits the bug: streaming off but + # interim assistant messages on, so prose arrives as commentary. + consumer.on_commentary("Now I get the full picture. Here's the situation.") + + # Run the drain loop concurrently while we block on the flush barrier + # from a worker thread (mirrors the agent thread calling the clarify + # callback while the consumer task drains on the event loop). + task = asyncio.create_task(consumer.run()) + flushed = await asyncio.to_thread(consumer.flush_pending_sync, 3.0) + + assert flushed is True, "flush_pending_sync should complete within timeout" + # The commentary must already be on screen by the time flush returns. + assert any( + "full picture" in content for _kind, content in sent_order + ), f"Commentary not delivered before flush returned: {sent_order!r}" + + consumer.finish() + await task + + @pytest.mark.asyncio + async def test_flush_times_out_when_consumer_not_running(self): + """If the consumer task is not running, flush_pending_sync returns + False (rather than hanging the caller forever).""" + adapter = MagicMock() + adapter.send = AsyncMock(return_value=SimpleNamespace(success=True, message_id="m1")) + adapter.edit_message = AsyncMock(return_value=SimpleNamespace(success=True)) + adapter.MAX_MESSAGE_LENGTH = 4096 + + config = StreamConsumerConfig(edit_interval=0.01, buffer_threshold=5) + consumer = GatewayStreamConsumer(adapter, "chat_123", config) + + # run() is never started — the barrier is never drained. + flushed = await asyncio.to_thread(consumer.flush_pending_sync, 0.1) + assert flushed is False + + @pytest.mark.asyncio + async def test_flush_with_empty_buffer_still_completes(self): + """A flush with nothing buffered still completes (no-op barrier).""" + adapter = MagicMock() + adapter.send = AsyncMock(return_value=SimpleNamespace(success=True, message_id="m1")) + adapter.edit_message = AsyncMock(return_value=SimpleNamespace(success=True)) + adapter.MAX_MESSAGE_LENGTH = 4096 + + config = StreamConsumerConfig(edit_interval=0.01, buffer_threshold=5) + consumer = GatewayStreamConsumer(adapter, "chat_123", config) + + task = asyncio.create_task(consumer.run()) + flushed = await asyncio.to_thread(consumer.flush_pending_sync, 3.0) + assert flushed is True + + consumer.finish() + await task + + @pytest.mark.asyncio + async def test_flush_completes_on_oversized_buffered_prose(self): + """Regression: oversized prose takes the overflow-split `continue` + path in run(), which previously skipped the flush-event set, stalling + the caller for the full timeout. The flush must still complete promptly + and the prose must be delivered before it returns.""" + adapter = MagicMock() + sent = [] + + async def _send(*args, **kwargs): + sent.append(kwargs.get("content", "")) + return SimpleNamespace(success=True, message_id=f"m{len(sent)}") + + adapter.send = AsyncMock(side_effect=_send) + adapter.edit_message = AsyncMock(return_value=SimpleNamespace(success=True)) + adapter.MAX_MESSAGE_LENGTH = 4096 + # Real splitter so the overflow branch (message_id is None + + # len > safe_limit) is actually taken. + adapter.truncate_message = ( + lambda text, limit, len_fn=len: [ + text[i:i + limit] for i in range(0, len(text), limit) + ] + ) + + config = StreamConsumerConfig(edit_interval=0.01, buffer_threshold=5) + consumer = GatewayStreamConsumer(adapter, "chat_123", config) + + # Commentary far larger than the platform limit → overflow split path. + big = "X" * 9000 + consumer.on_commentary(big) + + task = asyncio.create_task(consumer.run()) + # Tight timeout: if the continue-path skips the set, this returns False. + flushed = await asyncio.to_thread(consumer.flush_pending_sync, 2.0) + + assert flushed is True, ( + "flush stalled — overflow `continue` path did not signal the barrier" + ) + assert sent, "oversized prose was not delivered before flush returned" + assert sum(len(c) for c in sent) >= 9000 + + consumer.finish() + await task + + @pytest.mark.asyncio + async def test_flush_signaled_when_consumer_cancelled(self): + """If run() is cancelled while a flush barrier is queued, the finally + safety-net wakes the waiter rather than letting it hit the full + timeout.""" + adapter = MagicMock() + # Make the first send hang so the consumer is mid-iteration when we + # cancel it, with the flush barrier still in the queue behind it. + started = asyncio.Event() + + async def _slow_send(*args, **kwargs): + started.set() + await asyncio.sleep(60) + return SimpleNamespace(success=True, message_id="m1") + + adapter.send = AsyncMock(side_effect=_slow_send) + adapter.edit_message = AsyncMock(return_value=SimpleNamespace(success=True)) + adapter.MAX_MESSAGE_LENGTH = 4096 + + config = StreamConsumerConfig(edit_interval=0.01, buffer_threshold=5) + consumer = GatewayStreamConsumer(adapter, "chat_123", config) + consumer.on_commentary("first") + + task = asyncio.create_task(consumer.run()) + await started.wait() # consumer is now blocked inside the slow send + # Queue the flush barrier behind the hung send. + flush_done = asyncio.get_event_loop().run_in_executor( + None, consumer.flush_pending_sync, 5.0 + ) + await asyncio.sleep(0.05) + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + # The finally net should have drained + signaled the queued barrier. + flushed = await flush_done + assert flushed is True