ForcePilot/backend/package/yuxi/channel/extensions/slack/streaming.py

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)