141 lines
4.1 KiB
Python
141 lines
4.1 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import asyncio
|
||
|
|
import time
|
||
|
|
from dataclasses import dataclass
|
||
|
|
|
||
|
|
from slack_sdk.errors import SlackApiError
|
||
|
|
from slack_sdk.web.async_client import AsyncWebClient
|
||
|
|
|
||
|
|
from yuxi.channels.models import DeliveryResult
|
||
|
|
from yuxi.utils.logging_config import logger
|
||
|
|
|
||
|
|
DRAFT_STREAM_MAX_CHARS = 8000
|
||
|
|
DRAFT_STREAM_THROTTLE_MS = 1000
|
||
|
|
DRAFT_STREAM_MIN_THROTTLE_MS = 250
|
||
|
|
|
||
|
|
|
||
|
|
class DraftStream:
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
client: AsyncWebClient,
|
||
|
|
*,
|
||
|
|
max_chars: int = DRAFT_STREAM_MAX_CHARS,
|
||
|
|
throttle_ms: int = DRAFT_STREAM_THROTTLE_MS,
|
||
|
|
):
|
||
|
|
self._client = client
|
||
|
|
self._max_chars = max_chars
|
||
|
|
self._throttle_ms = max(throttle_ms, DRAFT_STREAM_MIN_THROTTLE_MS)
|
||
|
|
self._active: dict[str, _DraftState] = {}
|
||
|
|
self._lock = asyncio.Lock()
|
||
|
|
|
||
|
|
async def start(
|
||
|
|
self,
|
||
|
|
channel: str,
|
||
|
|
text: str,
|
||
|
|
*,
|
||
|
|
thread_ts: str | None = None,
|
||
|
|
) -> DeliveryResult:
|
||
|
|
try:
|
||
|
|
params: dict = {
|
||
|
|
"channel": channel,
|
||
|
|
"text": text[: self._max_chars],
|
||
|
|
"mrkdwn": True,
|
||
|
|
}
|
||
|
|
if thread_ts:
|
||
|
|
params["thread_ts"] = thread_ts
|
||
|
|
|
||
|
|
result = await self._client.chat_postMessage(**params)
|
||
|
|
ts = result.get("ts", "")
|
||
|
|
if not ts:
|
||
|
|
return DeliveryResult(success=False, error="No message ts returned")
|
||
|
|
|
||
|
|
async with self._lock:
|
||
|
|
self._active[channel] = _DraftState(
|
||
|
|
channel=channel,
|
||
|
|
ts=ts,
|
||
|
|
text=text,
|
||
|
|
last_update=time.monotonic(),
|
||
|
|
)
|
||
|
|
|
||
|
|
return DeliveryResult(success=True, message_id=ts)
|
||
|
|
|
||
|
|
except SlackApiError as e:
|
||
|
|
err = e.response.get("error", str(e))
|
||
|
|
logger.error(f"Draft stream start failed: {err}")
|
||
|
|
return DeliveryResult(success=False, error=err)
|
||
|
|
|
||
|
|
async def append(
|
||
|
|
self,
|
||
|
|
channel: str,
|
||
|
|
text: str,
|
||
|
|
*,
|
||
|
|
ts: str = "",
|
||
|
|
) -> DeliveryResult:
|
||
|
|
async with self._lock:
|
||
|
|
state = self._active.get(channel)
|
||
|
|
if not state:
|
||
|
|
return DeliveryResult(success=False, error="No active draft stream for channel")
|
||
|
|
|
||
|
|
now = time.monotonic()
|
||
|
|
elapsed_ms = (now - state.last_update) * 1000
|
||
|
|
if elapsed_ms < self._throttle_ms:
|
||
|
|
await asyncio.sleep((self._throttle_ms - elapsed_ms) / 1000)
|
||
|
|
|
||
|
|
try:
|
||
|
|
result = await self._client.chat_update(
|
||
|
|
channel=channel,
|
||
|
|
ts=ts or state.ts,
|
||
|
|
text=text[: self._max_chars],
|
||
|
|
mrkdwn=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
async with self._lock:
|
||
|
|
state.text = text
|
||
|
|
state.last_update = time.monotonic()
|
||
|
|
|
||
|
|
return DeliveryResult(
|
||
|
|
success=result.get("ok", False),
|
||
|
|
message_id=result.get("ts"),
|
||
|
|
error=result.get("error"),
|
||
|
|
)
|
||
|
|
|
||
|
|
except SlackApiError as e:
|
||
|
|
err = e.response.get("error", str(e))
|
||
|
|
if err == "message_not_found":
|
||
|
|
return await self.start(channel, text)
|
||
|
|
logger.error(f"Draft stream append failed: {err}")
|
||
|
|
return DeliveryResult(success=False, error=err)
|
||
|
|
|
||
|
|
async def seal(self, channel: str) -> DeliveryResult:
|
||
|
|
state = self._active.pop(channel, None)
|
||
|
|
if not state:
|
||
|
|
return DeliveryResult(success=False, error="No active draft stream")
|
||
|
|
return DeliveryResult(success=True, message_id=state.ts)
|
||
|
|
|
||
|
|
async def clear(self, channel: str) -> None:
|
||
|
|
self._active.pop(channel, None)
|
||
|
|
|
||
|
|
async def stop(self, channel: str) -> None:
|
||
|
|
self._active.pop(channel, None)
|
||
|
|
|
||
|
|
async def force_new_message(
|
||
|
|
self,
|
||
|
|
channel: str,
|
||
|
|
text: str,
|
||
|
|
) -> DeliveryResult:
|
||
|
|
self._active.pop(channel, None)
|
||
|
|
return await self.start(channel, text)
|
||
|
|
|
||
|
|
@property
|
||
|
|
def active_channels(self) -> set[str]:
|
||
|
|
return set(self._active.keys())
|
||
|
|
|
||
|
|
|
||
|
|
@dataclass
|
||
|
|
class _DraftState:
|
||
|
|
channel: str
|
||
|
|
ts: str
|
||
|
|
text: str
|
||
|
|
last_update: float
|