128 lines
3.5 KiB
Python
128 lines
3.5 KiB
Python
from __future__ import annotations
|
|
|
|
from typing import TYPE_CHECKING
|
|
|
|
from yuxi.channels.models import DeliveryResult
|
|
from yuxi.utils.logging_config import logger
|
|
|
|
from .formatter import markdown_to_matrix_html
|
|
|
|
if TYPE_CHECKING:
|
|
from nio import AsyncClient
|
|
|
|
|
|
class DraftStreamSession:
|
|
def __init__(self, room_id: str):
|
|
self.room_id = room_id
|
|
self.draft_event_id: str = ""
|
|
self.current_text: str = ""
|
|
self.is_active: bool = False
|
|
|
|
def start(self, event_id: str) -> None:
|
|
self.draft_event_id = event_id
|
|
self.current_text = ""
|
|
self.is_active = True
|
|
|
|
def update(self, text: str) -> None:
|
|
self.current_text = text
|
|
|
|
def finish(self) -> None:
|
|
self.is_active = False
|
|
|
|
|
|
_DRAFT_SESSIONS: dict[str, DraftStreamSession] = {}
|
|
|
|
|
|
def get_or_create_draft(room_id: str) -> DraftStreamSession:
|
|
if room_id not in _DRAFT_SESSIONS:
|
|
_DRAFT_SESSIONS[room_id] = DraftStreamSession(room_id)
|
|
return _DRAFT_SESSIONS[room_id]
|
|
|
|
|
|
def remove_draft(room_id: str) -> None:
|
|
_DRAFT_SESSIONS.pop(room_id, None)
|
|
|
|
|
|
async def send_draft_update(
|
|
client: AsyncClient,
|
|
room_id: str,
|
|
text: str,
|
|
draft_event_id: str | None = None,
|
|
finished: bool = False,
|
|
) -> DeliveryResult:
|
|
session = get_or_create_draft(room_id)
|
|
if draft_event_id and not session.draft_event_id:
|
|
session.start(draft_event_id)
|
|
|
|
if not session.is_active:
|
|
resp = await client.room_send(
|
|
room_id=room_id,
|
|
message_type="m.room.message",
|
|
content={
|
|
"msgtype": "m.text",
|
|
"body": text,
|
|
"format": "org.matrix.custom.html",
|
|
"formatted_body": markdown_to_matrix_html(text),
|
|
"org.matrix.msc4357.live": {
|
|
"draft": True,
|
|
"finished": finished,
|
|
},
|
|
},
|
|
)
|
|
session.start(resp.event_id)
|
|
return DeliveryResult(success=True, message_id=resp.event_id)
|
|
|
|
content = {
|
|
"msgtype": "m.text",
|
|
"body": text,
|
|
"format": "org.matrix.custom.html",
|
|
"formatted_body": markdown_to_matrix_html(text),
|
|
"m.new_content": {
|
|
"msgtype": "m.text",
|
|
"body": text,
|
|
"format": "org.matrix.custom.html",
|
|
"formatted_body": markdown_to_matrix_html(text),
|
|
},
|
|
"m.relates_to": {
|
|
"rel_type": "m.replace",
|
|
"event_id": session.draft_event_id,
|
|
},
|
|
}
|
|
|
|
if finished:
|
|
session.finish()
|
|
remove_draft(room_id)
|
|
|
|
try:
|
|
resp = await client.room_send(
|
|
room_id=room_id,
|
|
message_type="m.room.message",
|
|
content=content,
|
|
)
|
|
session.update(text)
|
|
return DeliveryResult(success=True, message_id=resp.event_id)
|
|
except Exception as e:
|
|
return DeliveryResult(success=False, error=str(e))
|
|
|
|
|
|
async def send_live_marker(
|
|
client: AsyncClient,
|
|
room_id: str,
|
|
event_id: str,
|
|
content: dict,
|
|
finished: bool = False,
|
|
) -> DeliveryResult:
|
|
content.setdefault("org.matrix.msc4357.live", {})
|
|
content["org.matrix.msc4357.live"]["finished"] = finished
|
|
|
|
try:
|
|
resp = await client.room_send(
|
|
room_id=room_id,
|
|
message_type="m.room.message",
|
|
content=content,
|
|
)
|
|
return DeliveryResult(success=True, message_id=resp.event_id)
|
|
except Exception as e:
|
|
logger.debug(f"Matrix MSC4357 live marker failed: {e}")
|
|
return DeliveryResult(success=False, error=str(e))
|