| """Gateway streaming consumer β bridges sync agent callbacks to async platform delivery. |
| |
| The agent fires stream_delta_callback(text) synchronously from its worker thread. |
| GatewayStreamConsumer: |
| 1. Receives deltas via on_delta() (thread-safe, sync) |
| 2. Queues them to an asyncio task via queue.Queue |
| 3. The async run() task buffers, rate-limits, and progressively edits |
| a single message on the target platform |
| |
| Design: Uses the edit transport (send initial message, then editMessageText). |
| This is universally supported across Telegram, Discord, and Slack. |
| |
| Credit: jobless0x (#774, #1312), OutThisLife (#798), clicksingh (#697). |
| """ |
|
|
| from __future__ import annotations |
|
|
| import asyncio |
| import logging |
| import queue |
| import re |
| import time |
| from dataclasses import dataclass |
| from typing import Any, Optional |
|
|
| logger = logging.getLogger("gateway.stream_consumer") |
|
|
| |
| _DONE = object() |
|
|
| |
| |
| _NEW_SEGMENT = object() |
|
|
| |
| |
| _COMMENTARY = object() |
|
|
|
|
| @dataclass |
| class StreamConsumerConfig: |
| """Runtime config for a single stream consumer instance.""" |
| edit_interval: float = 1.0 |
| buffer_threshold: int = 40 |
| cursor: str = " β" |
| buffer_only: bool = False |
|
|
|
|
| class GatewayStreamConsumer: |
| """Async consumer that progressively edits a platform message with streamed tokens. |
| |
| Usage:: |
| |
| consumer = GatewayStreamConsumer(adapter, chat_id, config, metadata=metadata) |
| # Pass consumer.on_delta as stream_delta_callback to AIAgent |
| agent = AIAgent(..., stream_delta_callback=consumer.on_delta) |
| # Start the consumer as an asyncio task |
| task = asyncio.create_task(consumer.run()) |
| # ... run agent in thread pool ... |
| consumer.finish() # signal completion |
| await task # wait for final edit |
| """ |
|
|
| |
| |
| _MAX_FLOOD_STRIKES = 3 |
|
|
| |
| |
| |
| _OPEN_THINK_TAGS = ( |
| "<REASONING_SCRATCHPAD>", "<think>", "<reasoning>", |
| "<THINKING>", "<thinking>", "<thought>", |
| ) |
| _CLOSE_THINK_TAGS = ( |
| "</REASONING_SCRATCHPAD>", "</think>", "</reasoning>", |
| "</THINKING>", "</thinking>", "</thought>", |
| ) |
|
|
| def __init__( |
| self, |
| adapter: Any, |
| chat_id: str, |
| config: Optional[StreamConsumerConfig] = None, |
| metadata: Optional[dict] = None, |
| ): |
| self.adapter = adapter |
| self.chat_id = chat_id |
| self.cfg = config or StreamConsumerConfig() |
| self.metadata = metadata |
| self._queue: queue.Queue = queue.Queue() |
| self._accumulated = "" |
| self._message_id: Optional[str] = None |
| self._already_sent = False |
| self._edit_supported = True |
| self._last_edit_time = 0.0 |
| self._last_sent_text = "" |
| self._fallback_final_send = False |
| self._fallback_prefix = "" |
| self._flood_strikes = 0 |
| self._current_edit_interval = self.cfg.edit_interval |
| self._final_response_sent = False |
| |
| |
| |
| |
| |
| self._adapter_requires_finalize: bool = ( |
| getattr(adapter, "REQUIRES_EDIT_FINALIZE", False) is True |
| ) |
|
|
| |
| self._in_think_block = False |
| self._think_buffer = "" |
|
|
| @property |
| def already_sent(self) -> bool: |
| """True if at least one message was sent or edited during the run.""" |
| return self._already_sent |
|
|
| @property |
| def final_response_sent(self) -> bool: |
| """True when the stream consumer delivered the final assistant reply.""" |
| return self._final_response_sent |
|
|
| def on_segment_break(self) -> None: |
| """Finalize the current stream segment and start a fresh message.""" |
| self._queue.put(_NEW_SEGMENT) |
|
|
| def on_commentary(self, text: str) -> None: |
| """Queue a completed interim assistant commentary message.""" |
| if text: |
| self._queue.put((_COMMENTARY, text)) |
|
|
| def _reset_segment_state(self, *, preserve_no_edit: bool = False) -> None: |
| if preserve_no_edit and self._message_id == "__no_edit__": |
| return |
| self._message_id = None |
| self._accumulated = "" |
| self._last_sent_text = "" |
| self._fallback_final_send = False |
| self._fallback_prefix = "" |
|
|
| def on_delta(self, text: str) -> None: |
| """Thread-safe callback β called from the agent's worker thread. |
| |
| When *text* is ``None``, signals a tool boundary: the current message |
| is finalized and subsequent text will be sent as a new message so it |
| appears below any tool-progress messages the gateway sent in between. |
| """ |
| if text: |
| self._queue.put(text) |
| elif text is None: |
| self.on_segment_break() |
|
|
| def finish(self) -> None: |
| """Signal that the stream is complete.""" |
| self._queue.put(_DONE) |
|
|
| |
| |
| |
| |
| |
| |
| |
|
|
| def _filter_and_accumulate(self, text: str) -> None: |
| """Add a text delta to the accumulated buffer, suppressing think blocks. |
| |
| Uses a state machine that tracks whether we are inside a |
| reasoning/thinking block. Text inside such blocks is silently |
| discarded. Partial tags at buffer boundaries are held back in |
| ``_think_buffer`` until enough characters arrive to decide. |
| """ |
| buf = self._think_buffer + text |
| self._think_buffer = "" |
|
|
| while buf: |
| if self._in_think_block: |
| |
| best_idx = -1 |
| best_len = 0 |
| for tag in self._CLOSE_THINK_TAGS: |
| idx = buf.find(tag) |
| if idx != -1 and (best_idx == -1 or idx < best_idx): |
| best_idx = idx |
| best_len = len(tag) |
|
|
| if best_len: |
| |
| self._in_think_block = False |
| buf = buf[best_idx + best_len:] |
| else: |
| |
| |
| max_tag = max(len(t) for t in self._CLOSE_THINK_TAGS) |
| self._think_buffer = buf[-max_tag:] if len(buf) > max_tag else buf |
| return |
| else: |
| |
| |
| |
| |
| best_idx = -1 |
| best_len = 0 |
| for tag in self._OPEN_THINK_TAGS: |
| search_start = 0 |
| while True: |
| idx = buf.find(tag, search_start) |
| if idx == -1: |
| break |
| |
| if idx == 0: |
| is_boundary = ( |
| not self._accumulated |
| or self._accumulated.endswith("\n") |
| ) |
| else: |
| preceding = buf[:idx] |
| last_nl = preceding.rfind("\n") |
| if last_nl == -1: |
| is_boundary = ( |
| (not self._accumulated |
| or self._accumulated.endswith("\n")) |
| and preceding.strip() == "" |
| ) |
| else: |
| is_boundary = preceding[last_nl + 1:].strip() == "" |
|
|
| if is_boundary and (best_idx == -1 or idx < best_idx): |
| best_idx = idx |
| best_len = len(tag) |
| break |
| search_start = idx + 1 |
|
|
| if best_len: |
| |
| self._accumulated += buf[:best_idx] |
| self._in_think_block = True |
| buf = buf[best_idx + best_len:] |
| else: |
| |
| held_back = 0 |
| for tag in self._OPEN_THINK_TAGS: |
| for i in range(1, len(tag)): |
| if buf.endswith(tag[:i]) and i > held_back: |
| held_back = i |
| if held_back: |
| self._accumulated += buf[:-held_back] |
| self._think_buffer = buf[-held_back:] |
| else: |
| self._accumulated += buf |
| return |
|
|
| def _flush_think_buffer(self) -> None: |
| """Flush any held-back partial-tag buffer into accumulated text. |
| |
| Called when the stream ends (got_done) so that partial text that |
| was held back waiting for a possible opening tag is not lost. |
| """ |
| if self._think_buffer and not self._in_think_block: |
| self._accumulated += self._think_buffer |
| self._think_buffer = "" |
|
|
| async def run(self) -> None: |
| """Async task that drains the queue and edits the platform message.""" |
| |
| _raw_limit = getattr(self.adapter, "MAX_MESSAGE_LENGTH", 4096) |
| _safe_limit = max(500, _raw_limit - len(self.cfg.cursor) - 100) |
|
|
| try: |
| while True: |
| |
| got_done = False |
| got_segment_break = False |
| commentary_text = None |
| while True: |
| try: |
| item = self._queue.get_nowait() |
| if item is _DONE: |
| got_done = True |
| break |
| if item is _NEW_SEGMENT: |
| got_segment_break = True |
| break |
| if isinstance(item, tuple) and len(item) == 2 and item[0] is _COMMENTARY: |
| commentary_text = item[1] |
| break |
| self._filter_and_accumulate(item) |
| except queue.Empty: |
| break |
|
|
| |
| |
| |
| if got_done: |
| self._flush_think_buffer() |
|
|
| |
| now = time.monotonic() |
| elapsed = now - self._last_edit_time |
| should_edit = ( |
| got_done |
| or got_segment_break |
| or commentary_text is not None |
| ) |
| if not self.cfg.buffer_only: |
| should_edit = should_edit or ( |
| (elapsed >= self._current_edit_interval |
| and self._accumulated) |
| or len(self._accumulated) >= self.cfg.buffer_threshold |
| ) |
|
|
| current_update_visible = False |
| if should_edit and self._accumulated: |
| |
| |
| if ( |
| len(self._accumulated) > _safe_limit |
| and self._message_id is None |
| ): |
| |
| |
| |
| |
| |
| chunks = self.adapter.truncate_message( |
| self._accumulated, _safe_limit |
| ) |
| for chunk in chunks: |
| await self._send_new_chunk(chunk, self._message_id) |
| self._accumulated = "" |
| self._last_sent_text = "" |
| self._last_edit_time = time.monotonic() |
| if got_done: |
| self._final_response_sent = self._already_sent |
| return |
| if got_segment_break: |
| self._message_id = None |
| self._fallback_final_send = False |
| self._fallback_prefix = "" |
| continue |
|
|
| |
| |
| while ( |
| len(self._accumulated) > _safe_limit |
| and self._message_id is not None |
| and self._edit_supported |
| ): |
| split_at = self._accumulated.rfind("\n", 0, _safe_limit) |
| if split_at < _safe_limit // 2: |
| split_at = _safe_limit |
| chunk = self._accumulated[:split_at] |
| ok = await self._send_or_edit(chunk) |
| if self._fallback_final_send or not ok: |
| |
| |
| |
| |
| |
| break |
| self._accumulated = self._accumulated[split_at:].lstrip("\n") |
| self._message_id = None |
| self._last_sent_text = "" |
|
|
| display_text = self._accumulated |
| if not got_done and not got_segment_break and commentary_text is None: |
| display_text += self.cfg.cursor |
|
|
| |
| |
| |
| |
| |
| |
| current_update_visible = await self._send_or_edit( |
| display_text, |
| finalize=got_segment_break, |
| ) |
| self._last_edit_time = time.monotonic() |
|
|
| if got_done: |
| |
| |
| |
| |
| if self._accumulated: |
| if self._fallback_final_send: |
| await self._send_fallback_final(self._accumulated) |
| elif ( |
| current_update_visible |
| and not self._adapter_requires_finalize |
| ): |
| |
| |
| |
| |
| self._final_response_sent = True |
| elif self._message_id: |
| |
| |
| |
| self._final_response_sent = await self._send_or_edit( |
| self._accumulated, finalize=True, |
| ) |
| elif not self._already_sent: |
| self._final_response_sent = await self._send_or_edit(self._accumulated) |
| return |
|
|
| if commentary_text is not None: |
| self._reset_segment_state() |
| await self._send_commentary(commentary_text) |
| self._last_edit_time = time.monotonic() |
| self._reset_segment_state() |
|
|
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| if got_segment_break: |
| |
| |
| |
| |
| |
| |
| |
| |
| if ( |
| self._accumulated |
| and not current_update_visible |
| and self._message_id |
| and self._message_id != "__no_edit__" |
| ): |
| await self._flush_segment_tail_on_edit_failure() |
| self._reset_segment_state(preserve_no_edit=True) |
|
|
| await asyncio.sleep(0.05) |
|
|
| except asyncio.CancelledError: |
| |
| _best_effort_ok = False |
| if self._accumulated and self._message_id: |
| try: |
| _best_effort_ok = bool(await self._send_or_edit(self._accumulated)) |
| except Exception: |
| pass |
| |
| |
| |
| |
| |
| |
| |
| if _best_effort_ok and not self._final_response_sent: |
| self._final_response_sent = True |
| except Exception as e: |
| logger.error("Stream consumer error: %s", e) |
|
|
| |
| |
| |
| _MEDIA_RE = re.compile(r'''[`"']?MEDIA:\s*\S+[`"']?''') |
|
|
| @staticmethod |
| def _clean_for_display(text: str) -> str: |
| """Strip MEDIA: directives and internal markers from text before display. |
| |
| The streaming path delivers raw text chunks that may include |
| ``MEDIA:<path>`` tags and ``[[audio_as_voice]]`` directives meant for |
| the platform adapter's post-processing. The actual media files are |
| delivered separately via ``_deliver_media_from_response()`` after the |
| stream finishes β we just need to hide the raw directives from the |
| user. |
| """ |
| if "MEDIA:" not in text and "[[audio_as_voice]]" not in text: |
| return text |
| cleaned = text.replace("[[audio_as_voice]]", "") |
| cleaned = GatewayStreamConsumer._MEDIA_RE.sub("", cleaned) |
| |
| cleaned = re.sub(r'\n{3,}', '\n\n', cleaned) |
| |
| return cleaned.rstrip() |
|
|
| async def _send_new_chunk(self, text: str, reply_to_id: Optional[str]) -> Optional[str]: |
| """Send a new message chunk, optionally threaded to a previous message. |
| |
| Returns the message_id so callers can thread subsequent chunks. |
| """ |
| text = self._clean_for_display(text) |
| if not text.strip(): |
| return reply_to_id |
| try: |
| meta = dict(self.metadata) if self.metadata else {} |
| result = await self.adapter.send( |
| chat_id=self.chat_id, |
| content=text, |
| reply_to=reply_to_id, |
| metadata=meta, |
| ) |
| if result.success and result.message_id: |
| self._message_id = str(result.message_id) |
| self._already_sent = True |
| self._last_sent_text = text |
| return str(result.message_id) |
| else: |
| self._edit_supported = False |
| return reply_to_id |
| except Exception as e: |
| logger.error("Stream send chunk error: %s", e) |
| return reply_to_id |
|
|
| def _visible_prefix(self) -> str: |
| """Return the visible text already shown in the streamed message.""" |
| prefix = self._last_sent_text or "" |
| if self.cfg.cursor and prefix.endswith(self.cfg.cursor): |
| prefix = prefix[:-len(self.cfg.cursor)] |
| return self._clean_for_display(prefix) |
|
|
| def _continuation_text(self, final_text: str) -> str: |
| """Return only the part of final_text the user has not already seen.""" |
| prefix = self._fallback_prefix or self._visible_prefix() |
| if prefix and final_text.startswith(prefix): |
| return final_text[len(prefix):].lstrip() |
| return final_text |
|
|
| @staticmethod |
| def _split_text_chunks(text: str, limit: int) -> list[str]: |
| """Split text into reasonably sized chunks for fallback sends.""" |
| if len(text) <= limit: |
| return [text] |
| chunks: list[str] = [] |
| remaining = text |
| while len(remaining) > limit: |
| split_at = remaining.rfind("\n", 0, limit) |
| if split_at < limit // 2: |
| split_at = limit |
| chunks.append(remaining[:split_at]) |
| remaining = remaining[split_at:].lstrip("\n") |
| if remaining: |
| chunks.append(remaining) |
| return chunks |
|
|
| async def _send_fallback_final(self, text: str) -> None: |
| """Send the final continuation after streaming edits stop working. |
| |
| Retries each chunk once on flood-control failures with a short delay. |
| """ |
| final_text = self._clean_for_display(text) |
| continuation = self._continuation_text(final_text) |
| self._fallback_final_send = False |
| if not continuation.strip(): |
| |
| |
| |
| |
| |
| |
| if final_text.strip() and final_text != self._visible_prefix(): |
| continuation = final_text |
| else: |
| |
| |
| |
| |
| |
| |
| |
| if ( |
| self._message_id |
| and self._last_sent_text |
| and self.cfg.cursor |
| and self._last_sent_text.endswith(self.cfg.cursor) |
| ): |
| clean_text = self._last_sent_text[:-len(self.cfg.cursor)] |
| try: |
| result = await self.adapter.edit_message( |
| chat_id=self.chat_id, |
| message_id=self._message_id, |
| content=clean_text, |
| ) |
| if result.success: |
| self._last_sent_text = clean_text |
| except Exception: |
| pass |
| self._already_sent = True |
| self._final_response_sent = True |
| return |
|
|
| raw_limit = getattr(self.adapter, "MAX_MESSAGE_LENGTH", 4096) |
| safe_limit = max(500, raw_limit - 100) |
| chunks = self._split_text_chunks(continuation, safe_limit) |
|
|
| last_message_id: Optional[str] = None |
| last_successful_chunk = "" |
| sent_any_chunk = False |
| for chunk in chunks: |
| |
| result = None |
| for attempt in range(2): |
| result = await self.adapter.send( |
| chat_id=self.chat_id, |
| content=chunk, |
| metadata=self.metadata, |
| ) |
| if result.success: |
| break |
| if attempt == 0 and self._is_flood_error(result): |
| logger.debug( |
| "Flood control on fallback send, retrying in 3s" |
| ) |
| await asyncio.sleep(3.0) |
| else: |
| break |
|
|
| if not result or not result.success: |
| if sent_any_chunk: |
| |
| |
| |
| self._already_sent = True |
| self._final_response_sent = True |
| self._message_id = last_message_id |
| self._last_sent_text = last_successful_chunk |
| self._fallback_prefix = "" |
| return |
| |
| |
| self._already_sent = False |
| self._message_id = None |
| self._last_sent_text = "" |
| self._fallback_prefix = "" |
| return |
| sent_any_chunk = True |
| last_successful_chunk = chunk |
| last_message_id = result.message_id or last_message_id |
|
|
| self._message_id = last_message_id |
| self._already_sent = True |
| self._final_response_sent = True |
| self._last_sent_text = chunks[-1] |
| self._fallback_prefix = "" |
|
|
| def _is_flood_error(self, result) -> bool: |
| """Check if a SendResult failure is due to flood control / rate limiting.""" |
| err = getattr(result, "error", "") or "" |
| err_lower = err.lower() |
| return "flood" in err_lower or "retry after" in err_lower or "rate" in err_lower |
|
|
| async def _flush_segment_tail_on_edit_failure(self) -> None: |
| """Deliver un-sent tail content before a segment-break reset. |
| |
| When an edit fails (flood control, transport error) and a tool |
| boundary arrives before the next retry, ``_accumulated`` holds text |
| that was generated but never shown to the user. Without this flush, |
| the segment reset would discard that tail and leave a frozen cursor |
| in the partial message. |
| |
| Sends the tail that sits after the last successfully-delivered |
| prefix as a new message, and best-effort strips the stuck cursor |
| from the previous partial message. |
| """ |
| if not self._fallback_final_send: |
| await self._try_strip_cursor() |
| visible = self._fallback_prefix or self._visible_prefix() |
| tail = self._accumulated |
| if visible and tail.startswith(visible): |
| tail = tail[len(visible):].lstrip() |
| tail = self._clean_for_display(tail) |
| if not tail.strip(): |
| return |
| try: |
| result = await self.adapter.send( |
| chat_id=self.chat_id, |
| content=tail, |
| metadata=self.metadata, |
| ) |
| if result.success: |
| self._already_sent = True |
| except Exception as e: |
| logger.error("Segment-break tail flush error: %s", e) |
|
|
| async def _try_strip_cursor(self) -> None: |
| """Best-effort edit to remove the cursor from the last visible message. |
| |
| Called when entering fallback mode so the user doesn't see a stuck |
| cursor (β) in the partial message. |
| """ |
| if not self._message_id or self._message_id == "__no_edit__": |
| return |
| prefix = self._visible_prefix() |
| if not prefix or not prefix.strip(): |
| return |
| try: |
| await self.adapter.edit_message( |
| chat_id=self.chat_id, |
| message_id=self._message_id, |
| content=prefix, |
| ) |
| self._last_sent_text = prefix |
| except Exception: |
| pass |
|
|
| async def _send_commentary(self, text: str) -> bool: |
| """Send a completed interim assistant commentary message.""" |
| text = self._clean_for_display(text) |
| if not text.strip(): |
| return False |
| try: |
| result = await self.adapter.send( |
| chat_id=self.chat_id, |
| content=text, |
| metadata=self.metadata, |
| ) |
| |
| |
| |
| |
| |
| return result.success |
| except Exception as e: |
| logger.error("Commentary send error: %s", e) |
| return False |
|
|
| async def _send_or_edit(self, text: str, *, finalize: bool = False) -> bool: |
| """Send or edit the streaming message. |
| |
| Returns True if the text was successfully delivered (sent or edited), |
| False otherwise. Callers like the overflow split loop use this to |
| decide whether to advance past the delivered chunk. |
| |
| ``finalize`` is True when this is the last edit in a streaming |
| sequence. |
| """ |
| |
| |
| |
| text = self._clean_for_display(text) |
| |
| |
| visible_without_cursor = text |
| if self.cfg.cursor: |
| visible_without_cursor = visible_without_cursor.replace(self.cfg.cursor, "") |
| _visible_stripped = visible_without_cursor.strip() |
| if not _visible_stripped: |
| return True |
| if not text.strip(): |
| return True |
| |
| |
| |
| |
| |
| |
| |
| |
| |
| _MIN_NEW_MSG_CHARS = 4 |
| if (self._message_id is None |
| and self.cfg.cursor |
| and self.cfg.cursor in text |
| and len(_visible_stripped) < _MIN_NEW_MSG_CHARS): |
| return True |
| try: |
| if self._message_id is not None: |
| if self._edit_supported: |
| |
| |
| |
| |
| |
| |
| if text == self._last_sent_text and not ( |
| finalize and self._adapter_requires_finalize |
| ): |
| return True |
| |
| result = await self.adapter.edit_message( |
| chat_id=self.chat_id, |
| message_id=self._message_id, |
| content=text, |
| finalize=finalize, |
| ) |
| if result.success: |
| self._already_sent = True |
| self._last_sent_text = text |
| |
| self._flood_strikes = 0 |
| return True |
| else: |
| |
| |
| |
| |
| if self._is_flood_error(result): |
| self._flood_strikes += 1 |
| self._current_edit_interval = min( |
| self._current_edit_interval * 2, 10.0, |
| ) |
| logger.debug( |
| "Flood control on edit (strike %d/%d), " |
| "backoff interval β %.1fs", |
| self._flood_strikes, |
| self._MAX_FLOOD_STRIKES, |
| self._current_edit_interval, |
| ) |
| if self._flood_strikes < self._MAX_FLOOD_STRIKES: |
| |
| |
| |
| self._last_edit_time = time.monotonic() |
| return False |
|
|
| |
| |
| |
| logger.debug( |
| "Edit failed (strikes=%d), entering fallback mode", |
| self._flood_strikes, |
| ) |
| self._fallback_prefix = self._visible_prefix() |
| self._fallback_final_send = True |
| self._edit_supported = False |
| self._already_sent = True |
| |
| |
| await self._try_strip_cursor() |
| return False |
| else: |
| |
| |
| return False |
| else: |
| |
| result = await self.adapter.send( |
| chat_id=self.chat_id, |
| content=text, |
| metadata=self.metadata, |
| ) |
| if result.success: |
| if result.message_id: |
| self._message_id = result.message_id |
| else: |
| self._edit_supported = False |
| self._already_sent = True |
| self._last_sent_text = text |
| if not result.message_id: |
| self._fallback_prefix = self._visible_prefix() |
| self._fallback_final_send = True |
| |
| |
| |
| self._message_id = "__no_edit__" |
| return True |
| else: |
| |
| self._edit_supported = False |
| return False |
| except Exception as e: |
| logger.error("Stream send/edit error: %s", e) |
| return False |
|
|