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

from __future__ import annotations

import logging
from dataclasses import dataclass

from yuxi.channel.config.defaults import (
    BLOCK_STREAMING_BREAK,
    BLOCK_STREAMING_CHUNK_BREAK_PREFERENCE,
    BLOCK_STREAMING_CHUNK_MAX_CHARS,
    BLOCK_STREAMING_CHUNK_MIN_CHARS,
    STREAMING_PREVIEW_MIN_INITIAL_CHARS,
    STREAMING_PREVIEW_THROTTLE_MS,
)

from .utils import get_nio

logger = logging.getLogger(__name__)


@dataclass
class MatrixDraftSession:
    target_id: str
    client: object
    replace_event_id: str | None = None
    streaming_mode: str = "partial"
    delivered: bool = False
    stopped: bool = False
    total_text: str = ""

    async def start(self, initial_text: str) -> str | None:
        content = {"msgtype": "m.text", "body": initial_text}
        resp = await self.client.room_send(self.target_id, "m.room.message", content)
        self.replace_event_id = resp.event_id
        self.delivered = True
        self.total_text = initial_text
        return resp.event_id

    async def append(self, text: str) -> None:
        if self.stopped or not self.replace_event_id:
            return
        self.total_text = text
        edit_content = {
            "msgtype": "m.text",
            "body": text,
            "m.new_content": {"msgtype": "m.text", "body": text},
            "m.relates_to": {
                "rel_type": "m.replace",
                "event_id": self.replace_event_id,
            },
        }
        try:
            await self.client.room_send(self.target_id, "m.room.message", edit_content)
        except Exception:
            logger.debug("Matrix draft append failed, falling back")
            self.stopped = True

    async def finalize(self, final_text: str = "") -> str:
        self.stopped = True
        if not self.replace_event_id:
            return ""
        text = final_text or self.total_text
        content = {"msgtype": "m.text", "body": text}
        edit_content = {
            **content,
            "m.new_content": content,
            "m.relates_to": {
                "rel_type": "m.replace",
                "event_id": self.replace_event_id,
            },
        }
        if self.streaming_mode == "quiet":
            edit_content["com.forcepilot.finalized_preview"] = True
        try:
            await self.client.room_send(self.target_id, "m.room.message", edit_content)
        except Exception:
            logger.warning("Matrix draft finalize failed for %s", self.replace_event_id)
        return self.replace_event_id

    async def abort(self) -> None:
        self.stopped = True
        if self.replace_event_id and self.total_text:
            await self.finalize(self.total_text)


streaming_mode = "partial"
preview_stream_throttle_ms = STREAMING_PREVIEW_THROTTLE_MS
preview_min_initial_chars = STREAMING_PREVIEW_MIN_INITIAL_CHARS
block_streaming_enabled = False
block_streaming_break = BLOCK_STREAMING_BREAK
block_streaming_chunk_min_chars = BLOCK_STREAMING_CHUNK_MIN_CHARS
block_streaming_chunk_max_chars = BLOCK_STREAMING_CHUNK_MAX_CHARS
block_streaming_chunk_break_preference = BLOCK_STREAMING_CHUNK_BREAK_PREFERENCE
block_streaming_coalesce_defaults = None


def create_draft_stream_session(
    target_id: str, config: dict = None, account_id: str = None, client: object = None
) -> object:
    if client is not None:
        return MatrixDraftSession(
            target_id=target_id,
            client=client,
            streaming_mode=_resolve_streaming_mode(config, account_id),
        )

    from .config import _apply_env_overrides, _dict_to_account

    cfg = config or {}
    aid = account_id or "default"
    account_data = cfg.get("accounts", {}).get(aid, {})
    account = _dict_to_account(account_data)
    account = _apply_env_overrides(account)

    if not account.homeserver or not account.access_token:
        logger.warning("Matrix draft stream session: account not configured")
        return None

    nio = get_nio()
    new_client = nio.AsyncClient(
        homeserver=account.homeserver,
        user=account.user_id,
        device_id=account.device_id,
    )
    new_client.access_token = account.access_token

    return MatrixDraftSession(
        target_id=target_id,
        client=new_client,
        streaming_mode=account.streaming or streaming_mode,
    )


def _resolve_streaming_mode(config: dict | None, account_id: str | None) -> str:
    cfg = config or {}
    aid = account_id or "default"
    account_data = cfg.get("accounts", {}).get(aid, {})
    return account_data.get("streaming", streaming_mode)


def create_block_chunker() -> object:
    from yuxi.channel.sdk import BlockChunker

    return BlockChunker()