import logging import time from dataclasses import dataclass, field from typing import Any from yuxi.channel.extensions.slack.constants import ( SLACK_STREAM_BUFFER_CHARS, SLACK_BENIGN_FINALIZE_CODES, SLACK_PREVIEW_THROTTLE_MS, SLACK_PREVIEW_MAX_CHARS, ) from yuxi.channel.extensions.slack.errors import SlackStreamNotDeliveredError logger = logging.getLogger(__name__) @dataclass class SlackStreamSession: client: Any = None channel: str = "" thread_ts: str = "" team_id: str | None = None user_id: str | None = None delivered: bool = False pending_text: str = "" stopped: bool = False _streamer: Any = field(default=None, repr=False) async def start(self, initial_text: str = "") -> None: kwargs: dict = {"channel": self.channel} if self.thread_ts: kwargs["thread_ts"] = self.thread_ts if self.team_id: kwargs["recipient_team_id"] = self.team_id if self.user_id: kwargs["recipient_user_id"] = self.user_id self._streamer = self.client.chat_stream(**kwargs) if initial_text: result = await self._streamer.append(markdown_text=initial_text) if result is not None: self.delivered = True async def append(self, text: str) -> None: if self.stopped or not text: return self.pending_text = text result = await self._streamer.append(markdown_text=text) if result is not None: self.delivered = True self.pending_text = "" async def stop(self, final_text: str = "") -> None: if self.stopped: return self.stopped = True try: from slack_sdk.errors import SlackApiError except ImportError: SlackApiError = None try: if final_text: await self._streamer.stop(markdown_text=final_text) else: await self._streamer.stop() self.delivered = True self.pending_text = "" except SlackApiError as e: error_code = e.response.get("error", "unknown") if error_code in SLACK_BENIGN_FINALIZE_CODES: if self.pending_text and not self.delivered: raise SlackStreamNotDeliveredError(self.pending_text, error_code) from e return raise @dataclass class SlackDraftStreamSession: client: Any = None channel: str = "" thread_ts: str = "" target_id: str = "" message_ts: str = "" content: str = "" throttle_ms: int = SLACK_PREVIEW_THROTTLE_MS max_chars: int = SLACK_PREVIEW_MAX_CHARS delivered: bool = False stopped: bool = False _last_flush: float = 0.0 async def start(self, initial_text: str = "") -> None: kwargs: dict = { "channel": self.target_id or self.channel, "text": initial_text or "...", "mrkdwn": True, } if self.thread_ts: kwargs["thread_ts"] = self.thread_ts result = await self.client.chat_postMessage(**kwargs) self.message_ts = result.get("ts", "") self.content = initial_text or "" self._last_flush = time.monotonic() async def append(self, text: str) -> None: if self.stopped or not text: return self.content += text now = time.monotonic() if now - self._last_flush < self.throttle_ms / 1000: return await self._flush() async def stop(self, final_text: str = "") -> None: if self.stopped: return self.stopped = True if final_text: self.content += final_text await self._flush() async def _flush(self) -> None: if not self.message_ts or not self.content: return try: truncated = self.content[: self.max_chars] await self.client.chat_update( channel=self.target_id or self.channel, ts=self.message_ts, text=truncated, mrkdwn=True, ) self.delivered = True self._last_flush = time.monotonic() except Exception as e: logger.warning("Draft stream update failed: %s", e) class SlackStreaming: streaming_mode = "native" preview_stream_throttle_ms = SLACK_PREVIEW_THROTTLE_MS preview_min_initial_chars = 18 block_streaming_enabled = False block_streaming_break = "text_end" block_streaming_chunk_min_chars = 800 block_streaming_chunk_max_chars = 1200 block_streaming_chunk_break_preference = "paragraph" block_streaming_coalesce_defaults = None def create_draft_stream_session(self, target_id: str) -> SlackDraftStreamSession: return SlackDraftStreamSession(target_id=target_id) def create_block_chunker(self) -> object: return {"mode": "length"} async def start_stream( self, client, channel: str, thread_ts: str = "", team_id: str | None = None, user_id: str | None = None, ) -> SlackStreamSession: session = SlackStreamSession( client=client, channel=channel, thread_ts=thread_ts, team_id=team_id, user_id=user_id, ) await session.start() return session async def append_stream(self, session: SlackStreamSession, text: str) -> None: await session.append(text) async def stop_stream(self, session: SlackStreamSession, final_text: str = "") -> None: await session.stop(final_text) def should_use_streaming(self, text_length: int) -> bool: return text_length > SLACK_STREAM_BUFFER_CHARS async def start_draft_stream( self, client, channel: str, thread_ts: str = "", initial_text: str = "", ) -> SlackDraftStreamSession: session = SlackDraftStreamSession( client=client, channel=channel, target_id=channel, thread_ts=thread_ts, ) await session.start(initial_text) return session async def append_draft_stream(self, session: SlackDraftStreamSession, text: str) -> None: await session.append(text) async def stop_draft_stream(self, session: SlackDraftStreamSession, final_text: str = "") -> None: await session.stop(final_text)