264 lines
8.0 KiB
Python
264 lines
8.0 KiB
Python
"""
|
||
测服自己的微盛监听名单与 last_seq。
|
||
|
||
本文件职责：把协同群 chat_id、按席/按房间的查询游标写到 REDIS_KEY_PREFIX 下。
|
||
禁止：写现网 ytd:prod: 前缀；禁止把游标当六态。
|
||
线程：可被建群线程与 Worker 轮询同时调用，只做单 key 读写。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import logging
|
||
import threading
|
||
import time
|
||
from typing import Optional
|
||
|
||
from agent.redis_coord import get_redis_runtime
|
||
from agent.redis_coord.client import RedisClient
|
||
from agent.redis_coord.keys import WSHOTO_LAST_SEQ, WSHOTO_ROOM_TICKET, WSHOTO_ROOMS
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
# 建群后或刚听到群话：立刻唤醒轮询，并在热窗内更勤地扫。
|
||
_poll_wake = threading.Event()
|
||
_hot_until: dict[str, float] = {}
|
||
_hot_lock = threading.Lock()
|
||
HOT_WINDOW_SEC = 90.0
|
||
|
||
|
||
def wake_wshoto_poll() -> None:
|
||
"""打断轮询等待。建群线程可调；真正拉消息仍在 Worker 轮询线程。"""
|
||
_poll_wake.set()
|
||
|
||
|
||
def wait_wshoto_poll(timeout: float) -> bool:
|
||
"""等下一轮或被唤醒。返回是否被踢醒。"""
|
||
fired = _poll_wake.wait(timeout=max(0.1, float(timeout)))
|
||
if fired:
|
||
_poll_wake.clear()
|
||
return fired
|
||
|
||
|
||
def mark_room_hot(chat_id: str, *, ttl_sec: float = HOT_WINDOW_SEC) -> None:
|
||
"""这群接下来几十秒更勤扫：刚建群或刚有人说话。"""
|
||
raw = (chat_id or "").strip()
|
||
if not raw:
|
||
return
|
||
until = time.time() + max(5.0, float(ttl_sec))
|
||
with _hot_lock:
|
||
_hot_until[raw] = until
|
||
wake_wshoto_poll()
|
||
|
||
|
||
def has_hot_room() -> bool:
|
||
"""是否还有热窗内的协同群。"""
|
||
return bool(list_hot_rooms())
|
||
|
||
|
||
def list_hot_rooms() -> list[str]:
|
||
"""当前热窗内的协同群（刚建群或刚有人说话）。"""
|
||
now = time.time()
|
||
with _hot_lock:
|
||
stale = [k for k, until in _hot_until.items() if until <= now]
|
||
for k in stale:
|
||
_hot_until.pop(k, None)
|
||
return [k for k, until in _hot_until.items() if until > now]
|
||
|
||
|
||
def is_room_hot(chat_id: str) -> bool:
|
||
"""这群是否还在热窗内。"""
|
||
raw = (chat_id or "").strip()
|
||
if not raw:
|
||
return False
|
||
now = time.time()
|
||
with _hot_lock:
|
||
until = _hot_until.get(raw) or 0.0
|
||
if until <= now:
|
||
_hot_until.pop(raw, None)
|
||
return False
|
||
return True
|
||
|
||
|
||
def plan_rooms_for_poll(
|
||
rooms: list[str],
|
||
*,
|
||
max_cold: int = 120,
|
||
redis: Optional[RedisClient] = None,
|
||
) -> list[str]:
|
||
"""
|
||
本轮要扫的房间：热群全部优先，冷群按游标轮转。
|
||
|
||
注意：这里只读游标、不推进。真正扫过多少冷群后，由 advance_poll_cursor 推进，
|
||
避免限频/超时提前停时把没扫到的群直接跳过。
|
||
"""
|
||
uniq: list[str] = []
|
||
seen: set[str] = set()
|
||
for raw in rooms:
|
||
room = (raw or "").strip()
|
||
if not room or room in seen:
|
||
continue
|
||
seen.add(room)
|
||
uniq.append(room)
|
||
hot = [r for r in uniq if is_room_hot(r)]
|
||
hot_set = set(hot)
|
||
cold = [r for r in uniq if r not in hot_set]
|
||
budget = max(0, int(max_cold))
|
||
if not cold or budget <= 0:
|
||
return hot
|
||
start = load_poll_cursor(redis=redis) % len(cold)
|
||
rotated = cold[start:] + cold[:start]
|
||
take = rotated[:budget]
|
||
# 热群在前，保证刚发过话/刚建的群本轮一定扫到。
|
||
return hot + [r for r in take if r not in hot_set]
|
||
|
||
|
||
def load_poll_cursor(*, redis: Optional[RedisClient] = None) -> int:
|
||
"""冷群轮转游标。"""
|
||
client = _client(redis)
|
||
if client is None:
|
||
return 0
|
||
raw_cur = client.raw.get(client.key(WSHOTO_LAST_SEQ, "poll_room_cursor"))
|
||
if isinstance(raw_cur, bytes):
|
||
raw_cur = raw_cur.decode("utf-8")
|
||
try:
|
||
return int(raw_cur or 0)
|
||
except (TypeError, ValueError):
|
||
return 0
|
||
|
||
|
||
def advance_poll_cursor(
|
||
cold_polled: int,
|
||
*,
|
||
redis: Optional[RedisClient] = None,
|
||
) -> None:
|
||
"""
|
||
按本轮实际扫过的冷群数推进游标。
|
||
|
||
限频或时间预算提前停时，只推进已扫到的，下轮从停住的地方继续。
|
||
"""
|
||
n = max(0, int(cold_polled))
|
||
if n <= 0:
|
||
return
|
||
client = _client(redis)
|
||
if client is None:
|
||
return
|
||
cur = load_poll_cursor(redis=client)
|
||
client.raw.set(client.key(WSHOTO_LAST_SEQ, "poll_room_cursor"), str(int(cur + n)))
|
||
|
||
|
||
def _client(redis: Optional[RedisClient] = None) -> Optional[RedisClient]:
|
||
if redis is not None:
|
||
return redis
|
||
try:
|
||
return get_redis_runtime().client
|
||
except Exception:
|
||
logger.exception("微盛监听名单：Redis 不可用")
|
||
return None
|
||
|
||
|
||
def watch_collab_room(
|
||
chat_id: str,
|
||
*,
|
||
work_order_no: str = "",
|
||
redis: Optional[RedisClient] = None,
|
||
) -> None:
|
||
"""
|
||
建群成功后登记 chat_id，供按房间补拉。
|
||
|
||
可同时记下房间→工单号，主账按群查找暂时不可用时也能对上单。
|
||
副作用：改 Redis。失败只打日志，不挡销售侧建群。
|
||
"""
|
||
raw = (chat_id or "").strip()
|
||
if not raw:
|
||
return
|
||
client = _client(redis)
|
||
if client is None:
|
||
return
|
||
key = client.key(WSHOTO_ROOMS)
|
||
current = client.raw.get(key) or ""
|
||
if isinstance(current, bytes):
|
||
current = current.decode("utf-8")
|
||
rooms = {x.strip() for x in str(current).split(",") if x.strip()}
|
||
rooms.add(raw)
|
||
client.raw.set(key, ",".join(sorted(rooms)))
|
||
no = (work_order_no or "").strip()
|
||
if no:
|
||
client.raw.set(client.key(WSHOTO_ROOM_TICKET, raw), no)
|
||
logger.info("存档监听已登记群数=%s", len(rooms))
|
||
mark_room_hot(raw)
|
||
|
||
|
||
def ticket_no_for_room(chat_id: str, *, redis: Optional[RedisClient] = None) -> str:
|
||
"""用协同群 chat_id 找回工单号（测服 Redis 短索引，不是六态）。"""
|
||
raw = (chat_id or "").strip()
|
||
if not raw:
|
||
return ""
|
||
client = _client(redis)
|
||
if client is None:
|
||
return ""
|
||
val = client.raw.get(client.key(WSHOTO_ROOM_TICKET, raw))
|
||
if val is None:
|
||
return ""
|
||
if isinstance(val, bytes):
|
||
val = val.decode("utf-8")
|
||
return str(val).strip()
|
||
|
||
|
||
def unwatch_collab_room(chat_id: str, *, redis: Optional[RedisClient] = None) -> None:
|
||
"""协同结束后去掉监听，避免解散群还被当成活群补拉。"""
|
||
raw = (chat_id or "").strip()
|
||
if not raw:
|
||
return
|
||
client = _client(redis)
|
||
if client is None:
|
||
return
|
||
key = client.key(WSHOTO_ROOMS)
|
||
current = client.raw.get(key) or ""
|
||
if isinstance(current, bytes):
|
||
current = current.decode("utf-8")
|
||
rooms = {x.strip() for x in str(current).split(",") if x.strip() and x.strip() != raw}
|
||
client.raw.set(key, ",".join(sorted(rooms)))
|
||
client.raw.delete(client.key(WSHOTO_ROOM_TICKET, raw))
|
||
|
||
|
||
def list_watched_rooms(*, redis: Optional[RedisClient] = None) -> list[str]:
|
||
"""当前要按房间补拉的协同群。"""
|
||
client = _client(redis)
|
||
if client is None:
|
||
return []
|
||
current = client.raw.get(client.key(WSHOTO_ROOMS)) or ""
|
||
if isinstance(current, bytes):
|
||
current = current.decode("utf-8")
|
||
return [x.strip() for x in str(current).split(",") if x.strip()]
|
||
|
||
|
||
def load_last_seq(scope: str, *, redis: Optional[RedisClient] = None) -> int:
|
||
"""
|
||
读取本前缀下的 last_seq。
|
||
|
||
scope：seat:AI001 或 room:wrxxx。没有记录返回 0（调用方改走时间窗，避免从 0 翻全量历史）。
|
||
"""
|
||
client = _client(redis)
|
||
if client is None:
|
||
return 0
|
||
val = client.raw.get(client.key(WSHOTO_LAST_SEQ, scope))
|
||
if val is None:
|
||
return 0
|
||
try:
|
||
return int(val)
|
||
except (TypeError, ValueError):
|
||
return 0
|
||
|
||
|
||
def save_last_seq(scope: str, seq: int, *, redis: Optional[RedisClient] = None) -> None:
|
||
"""推进测服自己的查询游标。seq 必须单调不减。"""
|
||
if seq <= 0:
|
||
return
|
||
client = _client(redis)
|
||
if client is None:
|
||
return
|
||
prev = load_last_seq(scope, redis=client)
|
||
if seq < prev:
|
||
return
|
||
client.raw.set(client.key(WSHOTO_LAST_SEQ, scope), str(int(seq)))
|