614 lines
22 KiB
Python
614 lines
22 KiB
Python
"""Restart-safe automatic replies over Huihui's self-hosted BOXIM."""
|
|
|
|
import asyncio
|
|
import hashlib
|
|
import logging
|
|
import re
|
|
import secrets
|
|
import time
|
|
from datetime import datetime, timedelta
|
|
from typing import Callable
|
|
|
|
from sqlalchemy.orm import Session
|
|
|
|
from models import (
|
|
Avatar,
|
|
TakeoverCursor,
|
|
TakeoverMessage,
|
|
TakeoverReplyTask,
|
|
User,
|
|
)
|
|
from services.boxim_client import BoxIMClient, BoxIMError
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
ACTIVE_TASK_STATUSES = ("pending", "generating", "ready", "sending")
|
|
GENERATABLE_TASK_STATUSES = ("pending",)
|
|
MAX_PROMPT_LENGTH = 4000
|
|
MAX_STALE_SECONDS = 120
|
|
STUCK_LOCK_SECONDS = 90
|
|
TAKEOVER_PERMISSION = "takeover"
|
|
|
|
|
|
def _utcnow() -> datetime:
|
|
return datetime.utcnow()
|
|
|
|
|
|
def _takeover_enabled(avatar: Avatar | None) -> bool:
|
|
if not avatar or avatar.status != "active":
|
|
return False
|
|
permissions = (avatar.config or {}).get("authorizationPermissions", [])
|
|
return isinstance(permissions, list) and TAKEOVER_PERMISSION in permissions
|
|
|
|
|
|
def _boxim_time(value, fallback: datetime) -> datetime:
|
|
try:
|
|
timestamp = float(value)
|
|
if timestamp > 10_000_000_000:
|
|
timestamp /= 1000
|
|
return datetime.utcfromtimestamp(timestamp)
|
|
except (TypeError, ValueError, OSError, OverflowError):
|
|
return fallback
|
|
|
|
|
|
def _numeric_id(value) -> int:
|
|
try:
|
|
return int(value)
|
|
except (TypeError, ValueError):
|
|
return 0
|
|
|
|
|
|
def _plain_text_reply(value: str) -> str:
|
|
"""BOXIM is plain text, so remove Markdown markers without damaging paragraphs."""
|
|
text = (value or "").replace("\r\n", "\n").replace("\r", "\n")
|
|
text = re.sub(r"```(?:\w+)?\n?(.*?)```", r"\1", text, flags=re.S)
|
|
text = re.sub(r"\*\*(.*?)\*\*|__(.*?)__", lambda m: m.group(1) or m.group(2), text)
|
|
text = re.sub(r"(?<!\*)\*([^*\n]+)\*(?!\*)", r"\1", text)
|
|
text = re.sub(r"`([^`]+)`", r"\1", text)
|
|
text = re.sub(r"^\s{0,3}#{1,6}\s*", "", text, flags=re.M)
|
|
lines = [line.strip() for line in text.split("\n")]
|
|
return "\n".join(line for line in lines if line).strip()
|
|
|
|
|
|
class TakeoverService:
|
|
"""Poll BOXIM, prepare replies during the grace period, then send at +3s."""
|
|
|
|
def __init__(
|
|
self,
|
|
session_factory: Callable[[], Session],
|
|
boxim_client: BoxIMClient,
|
|
*,
|
|
reply_delay_seconds: int = 3,
|
|
now: Callable[[], datetime] = _utcnow,
|
|
):
|
|
self.session_factory = session_factory
|
|
self.boxim = boxim_client
|
|
self.reply_delay_seconds = reply_delay_seconds
|
|
self.now = now
|
|
self._sessions: dict[str, dict] = {}
|
|
self._run_lock = asyncio.Lock()
|
|
|
|
async def poll_and_process_messages(self):
|
|
"""Run one complete cycle; polling always happens before reply dispatch."""
|
|
if self._run_lock.locked():
|
|
return
|
|
async with self._run_lock:
|
|
self._recover_stuck_tasks()
|
|
avatar_ids = self._enabled_avatar_ids()
|
|
self._cancel_disabled_tasks(set(avatar_ids))
|
|
for avatar_id in avatar_ids:
|
|
await self._sync_avatar(avatar_id)
|
|
|
|
generated = await self._prepare_replies()
|
|
if generated:
|
|
# Catch a human reply sent while the model was preparing its answer.
|
|
for avatar_id in avatar_ids:
|
|
await self._sync_avatar(avatar_id)
|
|
await self._dispatch_ready_replies()
|
|
|
|
def _enabled_avatar_ids(self) -> list[str]:
|
|
db = self.session_factory()
|
|
try:
|
|
return [
|
|
avatar.id
|
|
for avatar in db.query(Avatar).filter(Avatar.status == "active").all()
|
|
if _takeover_enabled(avatar)
|
|
]
|
|
finally:
|
|
db.close()
|
|
|
|
def _cancel_disabled_tasks(self, enabled_avatar_ids: set[str]):
|
|
db = self.session_factory()
|
|
try:
|
|
tasks = (
|
|
db.query(TakeoverReplyTask)
|
|
.filter(TakeoverReplyTask.status.in_(ACTIVE_TASK_STATUSES))
|
|
.all()
|
|
)
|
|
changed = False
|
|
for task in tasks:
|
|
if task.avatar_id not in enabled_avatar_ids:
|
|
task.status = "cancelled"
|
|
task.cancel_reason = "takeover_disabled"
|
|
task.locked_at = None
|
|
changed = True
|
|
if changed:
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
def _recover_stuck_tasks(self):
|
|
db = self.session_factory()
|
|
try:
|
|
threshold = self.now() - timedelta(seconds=STUCK_LOCK_SECONDS)
|
|
tasks = (
|
|
db.query(TakeoverReplyTask)
|
|
.filter(
|
|
TakeoverReplyTask.status.in_(("generating", "sending")),
|
|
TakeoverReplyTask.locked_at.isnot(None),
|
|
TakeoverReplyTask.locked_at < threshold,
|
|
)
|
|
.all()
|
|
)
|
|
for task in tasks:
|
|
task.status = "pending" if task.status == "generating" else "ready"
|
|
task.locked_at = None
|
|
task.last_error = "上次处理意外中断，已自动恢复"
|
|
if tasks:
|
|
db.commit()
|
|
finally:
|
|
db.close()
|
|
|
|
async def _boxim_session(self, user: User) -> dict:
|
|
token_fingerprint = hashlib.sha256((user.huihui_token or "").encode()).hexdigest()
|
|
cached = self._sessions.get(user.id)
|
|
if (
|
|
cached
|
|
and cached["expires_at"] > time.monotonic()
|
|
and cached["token_fingerprint"] == token_fingerprint
|
|
):
|
|
return cached
|
|
|
|
token_data = await self.boxim.exchange_access_token(user.huihui_token)
|
|
access_token = token_data["accessToken"]
|
|
profile = await self.boxim.get_self(access_token)
|
|
try:
|
|
expires_in = int(token_data.get("accessTokenExpiresIn") or 3600)
|
|
except (TypeError, ValueError):
|
|
expires_in = 3600
|
|
if expires_in > 86_400:
|
|
expires_in //= 1000
|
|
cache_for = max(60, min(expires_in - 60, 3600))
|
|
cached = {
|
|
"access_token": access_token,
|
|
"boxim_owner_id": str(profile["id"]),
|
|
"expires_at": time.monotonic() + cache_for,
|
|
"token_fingerprint": token_fingerprint,
|
|
}
|
|
self._sessions[user.id] = cached
|
|
return cached
|
|
|
|
def _forget_boxim_session(self, user_id: str):
|
|
self._sessions.pop(user_id, None)
|
|
|
|
def _disable_after_connection_failure(
|
|
self,
|
|
db: Session,
|
|
avatar: Avatar,
|
|
cursor: TakeoverCursor,
|
|
message: str,
|
|
):
|
|
permissions = (avatar.config or {}).get("authorizationPermissions", [])
|
|
avatar.config = {
|
|
**(avatar.config or {}),
|
|
"authorizationPermissions": [
|
|
permission
|
|
for permission in permissions
|
|
if permission != TAKEOVER_PERMISSION
|
|
],
|
|
}
|
|
cursor.last_error = message
|
|
cursor.last_polled_at = self.now()
|
|
tasks = (
|
|
db.query(TakeoverReplyTask)
|
|
.filter(
|
|
TakeoverReplyTask.avatar_id == avatar.id,
|
|
TakeoverReplyTask.status.in_(ACTIVE_TASK_STATUSES),
|
|
)
|
|
.all()
|
|
)
|
|
for task in tasks:
|
|
task.status = "cancelled"
|
|
task.cancel_reason = "connection_failed"
|
|
task.locked_at = None
|
|
|
|
async def _sync_avatar(self, avatar_id: str) -> bool:
|
|
db = self.session_factory()
|
|
try:
|
|
avatar = db.query(Avatar).filter(Avatar.id == avatar_id).first()
|
|
if not _takeover_enabled(avatar):
|
|
return False
|
|
user = db.query(User).filter(User.huihui_user_id == avatar.owner_id).first()
|
|
cursor = db.query(TakeoverCursor).filter(TakeoverCursor.avatar_id == avatar.id).first()
|
|
if not cursor:
|
|
cursor = TakeoverCursor(avatar_id=avatar.id, owner_id=avatar.owner_id)
|
|
db.add(cursor)
|
|
db.flush()
|
|
if not user or not user.huihui_token:
|
|
self._disable_after_connection_failure(
|
|
db,
|
|
avatar,
|
|
cursor,
|
|
"请重新登录会会生产账号后再开启主动接管",
|
|
)
|
|
db.commit()
|
|
return False
|
|
|
|
try:
|
|
session = await self._boxim_session(user)
|
|
owner_boxim_id = session["boxim_owner_id"]
|
|
if cursor.boxim_owner_id and cursor.boxim_owner_id != owner_boxim_id:
|
|
cursor.initialized = False
|
|
cursor.last_message_id = "0"
|
|
cursor.boxim_owner_id = owner_boxim_id
|
|
messages = await self.boxim.fetch_private_messages(
|
|
session["access_token"], cursor.last_message_id or "0"
|
|
)
|
|
except Exception as exc:
|
|
if isinstance(exc, BoxIMError) and exc.auth_error:
|
|
self._forget_boxim_session(user.id)
|
|
message = "BOXIM 授权已失效，请重新登录会会生产账号"
|
|
else:
|
|
message = f"BOXIM 暂时连接失败：{str(exc)[:160]}"
|
|
self._disable_after_connection_failure(db, avatar, cursor, message)
|
|
db.commit()
|
|
logger.warning("BOXIM sync failed for avatar %s: %s", avatar.id, exc)
|
|
return False
|
|
|
|
messages.sort(key=lambda item: (_numeric_id(item.get("id")), item.get("sendTime") or 0))
|
|
priming = not bool(cursor.initialized)
|
|
max_message_id = _numeric_id(cursor.last_message_id)
|
|
read_receipts: dict[str, int] = {}
|
|
for message in messages:
|
|
self._record_message(
|
|
db,
|
|
avatar,
|
|
cursor.boxim_owner_id,
|
|
message,
|
|
schedule_reply=not priming,
|
|
)
|
|
message_id = _numeric_id(message.get("id"))
|
|
max_message_id = max(max_message_id, message_id)
|
|
send_id = str(message.get("sendId") or "")
|
|
recv_id = str(message.get("recvId") or "")
|
|
if recv_id == cursor.boxim_owner_id and send_id and message_id:
|
|
read_receipts[send_id] = max(read_receipts.get(send_id, 0), message_id)
|
|
|
|
# BOXIM publishes this HTTP state change to connected socket clients.
|
|
# Do it before advancing the cursor so a failed receipt is retried.
|
|
for peer_id, message_id in read_receipts.items():
|
|
await self.boxim.mark_private_messages_read(
|
|
session["access_token"], peer_id, message_id
|
|
)
|
|
|
|
cursor.last_message_id = str(max_message_id)
|
|
cursor.initialized = True
|
|
cursor.last_polled_at = self.now()
|
|
cursor.last_error = ""
|
|
db.commit()
|
|
return True
|
|
except Exception:
|
|
db.rollback()
|
|
logger.exception("Failed to persist BOXIM messages for avatar %s", avatar_id)
|
|
return False
|
|
finally:
|
|
db.close()
|
|
|
|
def _record_message(
|
|
self,
|
|
db: Session,
|
|
avatar: Avatar,
|
|
boxim_owner_id: str,
|
|
message: dict,
|
|
*,
|
|
schedule_reply: bool,
|
|
):
|
|
message_id = str(message.get("id") or "").strip()
|
|
if not message_id:
|
|
return
|
|
local_id = str(message.get("localId") or "").strip() or None
|
|
if (
|
|
db.query(TakeoverMessage)
|
|
.filter(
|
|
TakeoverMessage.owner_id == avatar.owner_id,
|
|
TakeoverMessage.boxim_message_id == message_id,
|
|
)
|
|
.first()
|
|
):
|
|
return
|
|
|
|
send_id = str(message.get("sendId") or "")
|
|
recv_id = str(message.get("recvId") or "")
|
|
if send_id == boxim_owner_id:
|
|
direction, peer_id = "outgoing", recv_id
|
|
elif recv_id == boxim_owner_id:
|
|
direction, peer_id = "incoming", send_id
|
|
else:
|
|
return
|
|
if not peer_id:
|
|
return
|
|
|
|
now = self.now()
|
|
send_time = _boxim_time(message.get("sendTime"), now)
|
|
is_avatar = False
|
|
if direction == "outgoing" and local_id:
|
|
is_avatar = bool(
|
|
db.query(TakeoverReplyTask)
|
|
.filter(
|
|
TakeoverReplyTask.owner_id == avatar.owner_id,
|
|
TakeoverReplyTask.boxim_local_id == local_id,
|
|
TakeoverReplyTask.status == "sent",
|
|
)
|
|
.first()
|
|
)
|
|
|
|
event = TakeoverMessage(
|
|
avatar_id=avatar.id,
|
|
owner_id=avatar.owner_id,
|
|
boxim_message_id=message_id,
|
|
boxim_local_id=local_id,
|
|
peer_id=peer_id,
|
|
direction=direction,
|
|
message_type=int(message.get("type") or 0),
|
|
content=str(message.get("content") or ""),
|
|
is_avatar=is_avatar,
|
|
send_time=send_time,
|
|
)
|
|
db.add(event)
|
|
db.flush()
|
|
|
|
if direction == "outgoing":
|
|
if not is_avatar:
|
|
self._cancel_conversation(db, avatar.owner_id, peer_id, "owner_replied")
|
|
return
|
|
if not schedule_reply or event.message_type != 0 or not event.content.strip():
|
|
return
|
|
if (now - send_time).total_seconds() > MAX_STALE_SECONDS:
|
|
return
|
|
self._schedule_reply(db, avatar, event)
|
|
|
|
@staticmethod
|
|
def _cancel_conversation(db: Session, owner_id: str, peer_id: str, reason: str):
|
|
tasks = (
|
|
db.query(TakeoverReplyTask)
|
|
.filter(
|
|
TakeoverReplyTask.owner_id == owner_id,
|
|
TakeoverReplyTask.peer_id == peer_id,
|
|
TakeoverReplyTask.status.in_(ACTIVE_TASK_STATUSES),
|
|
)
|
|
.all()
|
|
)
|
|
for task in tasks:
|
|
task.status = "cancelled"
|
|
task.cancel_reason = reason
|
|
task.locked_at = None
|
|
|
|
def _schedule_reply(self, db: Session, avatar: Avatar, event: TakeoverMessage):
|
|
active_tasks = (
|
|
db.query(TakeoverReplyTask)
|
|
.filter(
|
|
TakeoverReplyTask.owner_id == avatar.owner_id,
|
|
TakeoverReplyTask.peer_id == event.peer_id,
|
|
TakeoverReplyTask.status.in_(("pending", "generating", "ready")),
|
|
)
|
|
.order_by(TakeoverReplyTask.created_at.desc())
|
|
.all()
|
|
)
|
|
prompt_parts = []
|
|
source_ids = []
|
|
if active_tasks:
|
|
latest = active_tasks[0]
|
|
prompt_parts.append(latest.prompt)
|
|
source_ids.extend(latest.source_message_ids or [])
|
|
for task in active_tasks:
|
|
task.status = "cancelled"
|
|
task.cancel_reason = "newer_incoming_message"
|
|
task.locked_at = None
|
|
prompt_parts.append(event.content.strip())
|
|
source_ids.append(event.boxim_message_id)
|
|
prompt = "\n".join(part for part in prompt_parts if part).strip()[-MAX_PROMPT_LENGTH:]
|
|
due_at = event.send_time + timedelta(seconds=self.reply_delay_seconds)
|
|
task_id = secrets.token_hex(16)
|
|
local_id = int(time.time() * 1000) * 1000 + secrets.randbelow(1000)
|
|
db.add(
|
|
TakeoverReplyTask(
|
|
id=task_id,
|
|
avatar_id=avatar.id,
|
|
owner_id=avatar.owner_id,
|
|
peer_id=event.peer_id,
|
|
trigger_message_id=event.boxim_message_id,
|
|
source_message_ids=source_ids,
|
|
prompt=prompt,
|
|
status="pending",
|
|
scheduled_at=due_at,
|
|
boxim_local_id=str(local_id),
|
|
)
|
|
)
|
|
|
|
async def _prepare_replies(self) -> int:
|
|
db = self.session_factory()
|
|
try:
|
|
task_ids = [
|
|
row[0]
|
|
for row in (
|
|
db.query(TakeoverReplyTask.id)
|
|
.filter(
|
|
TakeoverReplyTask.status.in_(GENERATABLE_TASK_STATUSES),
|
|
TakeoverReplyTask.response_text == "",
|
|
)
|
|
.order_by(TakeoverReplyTask.created_at.asc())
|
|
.limit(10)
|
|
.all()
|
|
)
|
|
]
|
|
finally:
|
|
db.close()
|
|
|
|
generated = 0
|
|
for task_id in task_ids:
|
|
if await asyncio.to_thread(self._generate_reply, task_id):
|
|
generated += 1
|
|
return generated
|
|
|
|
def _generate_reply(self, task_id: str) -> bool:
|
|
db = self.session_factory()
|
|
try:
|
|
task = db.query(TakeoverReplyTask).filter(TakeoverReplyTask.id == task_id).first()
|
|
if not task or task.status != "pending":
|
|
return False
|
|
avatar = db.query(Avatar).filter(Avatar.id == task.avatar_id).first()
|
|
if not _takeover_enabled(avatar):
|
|
task.status = "cancelled"
|
|
task.cancel_reason = "takeover_disabled"
|
|
db.commit()
|
|
return False
|
|
|
|
task.status = "generating"
|
|
task.locked_at = self.now()
|
|
db.commit()
|
|
|
|
excluded_ids = set(task.source_message_ids or [])
|
|
events = (
|
|
db.query(TakeoverMessage)
|
|
.filter(
|
|
TakeoverMessage.owner_id == task.owner_id,
|
|
TakeoverMessage.peer_id == task.peer_id,
|
|
)
|
|
.order_by(TakeoverMessage.send_time.desc())
|
|
.limit(30)
|
|
.all()
|
|
)
|
|
history = []
|
|
for event in reversed(events):
|
|
if event.boxim_message_id in excluded_ids or not event.content.strip():
|
|
continue
|
|
history.append(
|
|
{
|
|
"role": "user" if event.direction == "incoming" else "assistant",
|
|
"content": event.content.strip(),
|
|
}
|
|
)
|
|
history = history[-10:]
|
|
|
|
from routers.chat import _resolve_reply
|
|
|
|
result = _resolve_reply(db, avatar, task.prompt, history)
|
|
answer = _plain_text_reply(result.get("answer", ""))
|
|
db.refresh(task)
|
|
if task.status != "generating":
|
|
return False
|
|
if not answer:
|
|
raise RuntimeError("分身没有生成有效回复")
|
|
task.response_text = answer
|
|
task.status = "ready"
|
|
task.locked_at = None
|
|
task.last_error = ""
|
|
db.commit()
|
|
return True
|
|
except Exception as exc:
|
|
db.rollback()
|
|
task = db.query(TakeoverReplyTask).filter(TakeoverReplyTask.id == task_id).first()
|
|
if task and task.status in ("pending", "generating"):
|
|
task.attempts = (task.attempts or 0) + 1
|
|
task.status = "pending" if task.attempts < 3 else "failed"
|
|
task.locked_at = None
|
|
task.last_error = str(exc)[:300]
|
|
db.commit()
|
|
logger.warning("Failed to prepare takeover reply %s: %s", task_id, exc)
|
|
return False
|
|
finally:
|
|
db.close()
|
|
|
|
async def _dispatch_ready_replies(self):
|
|
db = self.session_factory()
|
|
try:
|
|
task_ids = [
|
|
row[0]
|
|
for row in (
|
|
db.query(TakeoverReplyTask.id)
|
|
.filter(
|
|
TakeoverReplyTask.status == "ready",
|
|
TakeoverReplyTask.scheduled_at <= self.now(),
|
|
)
|
|
.order_by(TakeoverReplyTask.scheduled_at.asc())
|
|
.limit(10)
|
|
.all()
|
|
)
|
|
]
|
|
finally:
|
|
db.close()
|
|
|
|
for task_id in task_ids:
|
|
await self._send_task(task_id)
|
|
|
|
async def _send_task(self, task_id: str) -> bool:
|
|
db = self.session_factory()
|
|
user = None
|
|
try:
|
|
task = db.query(TakeoverReplyTask).filter(TakeoverReplyTask.id == task_id).first()
|
|
if not task or task.status != "ready":
|
|
return False
|
|
avatar = db.query(Avatar).filter(Avatar.id == task.avatar_id).first()
|
|
if not _takeover_enabled(avatar):
|
|
task.status = "cancelled"
|
|
task.cancel_reason = "takeover_disabled"
|
|
db.commit()
|
|
return False
|
|
if (self.now() - task.scheduled_at).total_seconds() > MAX_STALE_SECONDS:
|
|
task.status = "cancelled"
|
|
task.cancel_reason = "stale_reply"
|
|
db.commit()
|
|
return False
|
|
user = db.query(User).filter(User.huihui_user_id == task.owner_id).first()
|
|
if not user or not user.huihui_token:
|
|
raise BoxIMError("缺少会会登录凭证", auth_error=True)
|
|
|
|
task.status = "sending"
|
|
task.locked_at = self.now()
|
|
db.commit()
|
|
session = await self._boxim_session(user)
|
|
result = await self.boxim.send_private_message(
|
|
session["access_token"],
|
|
task.peer_id,
|
|
task.response_text,
|
|
local_id=task.boxim_local_id,
|
|
)
|
|
db.refresh(task)
|
|
if task.status != "sending":
|
|
return False
|
|
task.status = "sent"
|
|
task.sent_at = self.now()
|
|
task.locked_at = None
|
|
task.last_error = ""
|
|
task.boxim_sent_message_id = str(result.get("id") or "")
|
|
db.commit()
|
|
logger.info("BOXIM takeover reply sent for task %s", task.id)
|
|
return True
|
|
except Exception as exc:
|
|
db.rollback()
|
|
if user and isinstance(exc, BoxIMError) and exc.auth_error:
|
|
self._forget_boxim_session(user.id)
|
|
task = db.query(TakeoverReplyTask).filter(TakeoverReplyTask.id == task_id).first()
|
|
if task and task.status in ("ready", "sending"):
|
|
task.attempts = (task.attempts or 0) + 1
|
|
task.status = "ready" if task.attempts < 3 else "failed"
|
|
task.locked_at = None
|
|
task.last_error = str(exc)[:300]
|
|
if task.status == "ready":
|
|
task.scheduled_at = self.now() + timedelta(seconds=2 ** task.attempts)
|
|
db.commit()
|
|
logger.warning("Failed to send takeover reply %s: %s", task_id, exc)
|
|
return False
|
|
finally:
|
|
db.close()
|