Kris
ef5483dc1a
主要变更： 1. 修复速率限流器使用setdefault替代重复创建令牌桶 2. 重构交互注册表匹配逻辑，优化精确匹配查找 3. 重构去重缓存逻辑，移到适配器实例方法 4. 重构发送URL解析，增加合法性校验并拆分公共方法 5. 优化流式消息处理逻辑，简化flush_controller调用 6. 重构群聊类型判断代码，简化语法 7. 修复重连管理器对None类型关闭分类的处理 8. 新增消息缓存、线程模拟器、发送初始化模块 9. 重构凭证备份与会话存储逻辑，支持环境变量指定状态目录 10. 新增配置提示与向导二维码绑定功能 11. 优化媒体上传逻辑，增加重试机制与缓存 12. 新增审批键盘模板构建函数 13. 重构消息格式处理，修正媒体发送字段与长度限制 14. 修复令牌过期时间计算，使用time.time替代monotonic 15. 新增群组激活缓冲区与用户追踪器增强功能 16. 修复换行符问题，统一文件结尾格式
377 lines
13 KiB
Python
377 lines
13 KiB
Python
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from collections.abc import AsyncGenerator
|
|
from dataclasses import dataclass, field
|
|
from enum import Enum, auto
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class StreamEventType(Enum):
|
|
FLUSH = auto()
|
|
COMPLETE = auto()
|
|
ERROR = auto()
|
|
|
|
|
|
class FlushStrategy(Enum):
|
|
PER_CHAR = auto()
|
|
INTERVAL = auto()
|
|
BACKLOG = auto()
|
|
|
|
|
|
@dataclass
|
|
class StreamEvent:
|
|
event_type: StreamEventType
|
|
content: str = ""
|
|
full_content: str = ""
|
|
error: str | None = None
|
|
|
|
|
|
@dataclass
|
|
class StreamCheckpoint:
|
|
chat_id: str
|
|
msg_id: str = ""
|
|
content: str = ""
|
|
sent_seq: int = 0
|
|
timestamp: float = field(default_factory=time.time)
|
|
|
|
@property
|
|
def is_stale(self) -> bool:
|
|
return time.time() - self.timestamp > 600
|
|
|
|
|
|
BATCH_AFTER_GAP = 1.5
|
|
|
|
|
|
@dataclass
|
|
class FlushController:
|
|
strategy: FlushStrategy = FlushStrategy.BACKLOG
|
|
flush_interval: float = 0.6
|
|
backlog_threshold: int = 30
|
|
max_retries: int = 3
|
|
retry_delay: float = 1.0
|
|
seq_counter: int = 0
|
|
_buffer: list[str] = field(default_factory=list, repr=False)
|
|
_last_flush: float = 0.0
|
|
_accumulated: str = ""
|
|
_flush_in_progress: bool = False
|
|
_needs_reflush: bool = False
|
|
_last_chunk_time: float = 0.0
|
|
|
|
def reset(self) -> None:
|
|
self._buffer = []
|
|
self._last_flush = 0.0
|
|
self._accumulated = ""
|
|
self.seq_counter = 0
|
|
self._flush_in_progress = False
|
|
self._needs_reflush = False
|
|
self._last_chunk_time = 0.0
|
|
|
|
def feed(self, chunk: str) -> list[str]:
|
|
self._buffer.append(chunk)
|
|
self._accumulated += chunk
|
|
self._last_chunk_time = time.monotonic()
|
|
|
|
if self.strategy == FlushStrategy.PER_CHAR:
|
|
return self._flush_all()
|
|
|
|
if self.strategy == FlushStrategy.BACKLOG and len(self._accumulated) >= self.backlog_threshold:
|
|
return self._flush_all()
|
|
|
|
now = time.monotonic()
|
|
if self.strategy == FlushStrategy.INTERVAL and now - self._last_flush >= self.flush_interval:
|
|
return self._flush_all()
|
|
|
|
if self._last_chunk_time > 0 and now - self._last_chunk_time >= BATCH_AFTER_GAP:
|
|
return self._flush_all()
|
|
|
|
return []
|
|
|
|
def flush_remaining(self) -> tuple[list[str], str]:
|
|
if self.strategy == FlushStrategy.PER_CHAR:
|
|
return [], self._accumulated
|
|
batches = self._flush_all() if self._buffer else []
|
|
return batches, self._accumulated
|
|
|
|
def mark_flush_in_progress(self) -> bool:
|
|
if self._flush_in_progress:
|
|
self._needs_reflush = True
|
|
return False
|
|
self._flush_in_progress = True
|
|
return True
|
|
|
|
def mark_flush_done(self) -> bool:
|
|
self._flush_in_progress = False
|
|
needs = self._needs_reflush
|
|
self._needs_reflush = False
|
|
return needs
|
|
|
|
def _flush_all(self) -> list[str]:
|
|
if not self._buffer:
|
|
return []
|
|
batches = list(self._buffer)
|
|
self._buffer = []
|
|
self._accumulated = ""
|
|
self._last_flush = time.monotonic()
|
|
return batches
|
|
|
|
|
|
class C2CStreamingController:
|
|
def __init__(
|
|
self,
|
|
send_message_fn,
|
|
retry_delay: float = 1.0,
|
|
max_retries: int = 3,
|
|
flush_interval: float = 0.6,
|
|
) -> None:
|
|
self._send_message_fn = send_message_fn
|
|
self._retry_delay = retry_delay
|
|
self._max_retries = max_retries
|
|
self._flush_interval = flush_interval
|
|
self._flush_controller = FlushController(strategy=FlushStrategy.BACKLOG, flush_interval=flush_interval)
|
|
self._any_chunk_delivered: bool = False
|
|
self._static_fallback_msg: str = ""
|
|
self._active_checkpoints: dict[str, StreamCheckpoint] = {}
|
|
|
|
@property
|
|
def flush_controller(self) -> FlushController:
|
|
return self._flush_controller
|
|
|
|
async def stream(
|
|
self,
|
|
chat_id: str,
|
|
msg_id: str,
|
|
content_generator: AsyncGenerator[str, None],
|
|
event_id: str = "",
|
|
) -> bool:
|
|
self._flush_controller.reset()
|
|
self._any_chunk_delivered = False
|
|
self._static_fallback_msg = ""
|
|
|
|
checkpoint = StreamCheckpoint(chat_id=chat_id, msg_id=msg_id)
|
|
self._active_checkpoints[chat_id] = checkpoint
|
|
|
|
collected = ""
|
|
in_media_interrupt = False
|
|
media_buffer = ""
|
|
|
|
try:
|
|
async for chunk in content_generator:
|
|
if not chunk:
|
|
continue
|
|
|
|
media_splits = self._split_on_media_tags(chunk)
|
|
for segment, is_media in media_splits:
|
|
if is_media:
|
|
in_media_interrupt = True
|
|
media_buffer += segment
|
|
else:
|
|
if in_media_interrupt and media_buffer:
|
|
checkpoint.content = collected
|
|
await self._flush_and_end_stream(chat_id, msg_id, checkpoint, event_id)
|
|
await self._send_media_interruption(chat_id, media_buffer)
|
|
media_buffer = ""
|
|
in_media_interrupt = False
|
|
|
|
new_msg_id = f"{msg_id}_resume_{checkpoint.sent_seq}"
|
|
checkpoint.msg_id = new_msg_id
|
|
self._flush_controller.reset()
|
|
|
|
collected += segment
|
|
batches = self._flush_controller.feed(segment)
|
|
for batch in batches:
|
|
if not self._flush_controller.mark_flush_in_progress():
|
|
continue
|
|
self._flush_controller.seq_counter += 1
|
|
success = await self._send_stream_chunk(
|
|
chat_id,
|
|
checkpoint.msg_id or msg_id,
|
|
batch,
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
)
|
|
_needs_reflush = self._flush_controller.mark_flush_done()
|
|
if success:
|
|
self._any_chunk_delivered = True
|
|
else:
|
|
logger.warning(
|
|
"C2CStreaming: flush failed for seq %d", self._flush_controller.seq_counter
|
|
)
|
|
|
|
self._static_fallback_msg = collected
|
|
|
|
flush_batches, full = self._flush_controller.flush_remaining()
|
|
for batch in flush_batches:
|
|
if not self._flush_controller.mark_flush_in_progress():
|
|
continue
|
|
self._flush_controller.seq_counter += 1
|
|
success = await self._send_stream_chunk(
|
|
chat_id,
|
|
checkpoint.msg_id or msg_id,
|
|
batch,
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
)
|
|
_needs_reflush = self._flush_controller.mark_flush_done()
|
|
if success:
|
|
self._any_chunk_delivered = True
|
|
|
|
self._flush_controller.seq_counter += 1
|
|
success = await self._send_stream_chunk(
|
|
chat_id,
|
|
checkpoint.msg_id or msg_id,
|
|
"",
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
is_end=True,
|
|
)
|
|
self._active_checkpoints.pop(chat_id, None)
|
|
return success
|
|
|
|
except Exception as e:
|
|
logger.exception("C2CStreaming: stream error for %s", chat_id)
|
|
self._active_checkpoints.pop(chat_id, None)
|
|
return await self._cancel_stream(chat_id, msg_id, event_id, str(e))
|
|
|
|
async def _send_stream_chunk(
|
|
self,
|
|
chat_id: str,
|
|
msg_id: str,
|
|
content: str,
|
|
msg_seq: int,
|
|
event_id: str = "",
|
|
is_end: bool = False,
|
|
) -> bool:
|
|
for attempt in range(self._max_retries):
|
|
try:
|
|
payload = {
|
|
"content": content,
|
|
"msg_type": 0,
|
|
"msg_id": msg_id,
|
|
"msg_seq": msg_seq,
|
|
"stream": {"state": 2 if is_end else 1},
|
|
}
|
|
if event_id:
|
|
payload["event_id"] = event_id
|
|
|
|
response = await self._send_message_fn(chat_id, payload)
|
|
if response is not None:
|
|
cp = self._active_checkpoints.get(chat_id)
|
|
if cp:
|
|
cp.sent_seq = msg_seq
|
|
cp.timestamp = time.time()
|
|
return response is not None
|
|
except Exception:
|
|
if attempt < self._max_retries - 1:
|
|
await asyncio.sleep(self._retry_delay)
|
|
else:
|
|
logger.exception("C2CStreaming: send chunk failed after %d retries", self._max_retries)
|
|
return False
|
|
|
|
@staticmethod
|
|
def _split_on_media_tags(text: str) -> list[tuple[str, bool]]:
|
|
import re
|
|
|
|
from .media_tags import _IMG_TAG_RE, _MEDIA_TAG_RE, _VIDEO_TAG_RE
|
|
|
|
combined_re = re.compile(f"({_MEDIA_TAG_RE.pattern}|{_IMG_TAG_RE.pattern}|{_VIDEO_TAG_RE.pattern})")
|
|
|
|
results: list[tuple[str, bool]] = []
|
|
last_end = 0
|
|
for match in combined_re.finditer(text):
|
|
if match.start() > last_end:
|
|
results.append((text[last_end : match.start()], False))
|
|
results.append((match.group(0), True))
|
|
last_end = match.end()
|
|
|
|
if last_end < len(text):
|
|
results.append((text[last_end:], False))
|
|
|
|
return results
|
|
|
|
async def _flush_and_end_stream(
|
|
self, chat_id: str, msg_id: str, checkpoint: StreamCheckpoint, event_id: str
|
|
) -> None:
|
|
flush_batches, _ = self._flush_controller.flush_remaining()
|
|
for batch in flush_batches:
|
|
self._flush_controller.seq_counter += 1
|
|
await self._send_stream_chunk(
|
|
chat_id,
|
|
msg_id,
|
|
batch,
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
)
|
|
|
|
self._flush_controller.seq_counter += 1
|
|
await self._send_stream_chunk(
|
|
chat_id,
|
|
msg_id,
|
|
"",
|
|
self._flush_controller.seq_counter,
|
|
event_id,
|
|
is_end=True,
|
|
)
|
|
|
|
async def _send_media_interruption(self, chat_id: str, media_tags: str) -> None:
|
|
try:
|
|
payload = {
|
|
"content": media_tags,
|
|
"msg_type": 0,
|
|
"msg_id": "",
|
|
"msg_seq": 0,
|
|
}
|
|
await self._send_message_fn(chat_id, payload)
|
|
except Exception:
|
|
logger.exception("C2CStreaming: media interruption send failed")
|
|
|
|
async def _cancel_stream(self, chat_id: str, msg_id: str, event_id: str, reason: str) -> bool:
|
|
delivered_fallback = False
|
|
if self._static_fallback_msg and not self._any_chunk_delivered:
|
|
try:
|
|
static_payload = {
|
|
"content": self._static_fallback_msg[:2000],
|
|
"msg_type": 0,
|
|
"msg_id": msg_id,
|
|
"msg_seq": 0,
|
|
}
|
|
if event_id:
|
|
static_payload["event_id"] = event_id
|
|
await self._send_message_fn(chat_id, static_payload)
|
|
delivered_fallback = True
|
|
logger.info("C2CStreaming: delivered static fallback message for %s", chat_id)
|
|
except Exception:
|
|
logger.exception("C2CStreaming: static fallback delivery failed for %s", chat_id)
|
|
|
|
if not delivered_fallback:
|
|
try:
|
|
payload = {
|
|
"content": reason,
|
|
"msg_type": 0,
|
|
"msg_id": msg_id,
|
|
"msg_seq": 0,
|
|
"stream": {"state": 0},
|
|
}
|
|
if event_id:
|
|
payload["event_id"] = event_id
|
|
await self._send_message_fn(chat_id, payload)
|
|
except Exception:
|
|
pass
|
|
return False
|
|
|
|
def get_checkpoint(self, chat_id: str) -> StreamCheckpoint | None:
|
|
return self._active_checkpoints.get(chat_id)
|
|
|
|
def has_pending_stream(self, chat_id: str) -> bool:
|
|
cp = self._active_checkpoints.get(chat_id)
|
|
return cp is not None and not cp.is_stale
|
|
|
|
def cleanup_stale_checkpoints(self) -> int:
|
|
stale = [cid for cid, cp in self._active_checkpoints.items() if cp.is_stale]
|
|
for cid in stale:
|
|
self._active_checkpoints.pop(cid, None)
|
|
return len(stale)
|