109 lines
4.8 KiB
Python
109 lines
4.8 KiB
Python
"""Per-chat FIFO send queues with token-bucket rate limiting for WeCom.
|
|
|
|
Mirrors OpenClaw's chat-queue.ts (serial per chat) plus a token bucket that
|
|
keeps each chat under WeCom's 30 msgs/min/chat limit (errcode 846607). Two
|
|
lanes per chat: a normal lane and a high-priority control lane (approval
|
|
prompts, finalize frames, error notices) backed by a reserved token pool.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import logging
|
|
import time
|
|
from typing import Dict
|
|
|
|
logger = logging.getLogger("plugins.platforms.wecom.adapter")
|
|
|
|
|
|
class ChatSendQueueMixin:
|
|
"""Expects ``_chat_queues/_chat_workers/_control_queues/_control_workers/_chat_token_usage`` dicts."""
|
|
|
|
# Token bucket: 30 tokens/min per chat, split between normal and reserved (control) quota.
|
|
_BUCKET_MAX_TOKENS = 30
|
|
_BUCKET_NORMAL_TOKENS = 24
|
|
_BUCKET_RESERVED_TOKENS = 6
|
|
|
|
def _get_token_usage(self, chat_id: str) -> Dict[str, float]:
|
|
"""Get or create token usage tracking for a chat."""
|
|
key = str(chat_id or "").strip()
|
|
if key not in self._chat_token_usage:
|
|
self._chat_token_usage[key] = {"normal": 0.0, "reserved": 0.0, "last_reset": time.monotonic()}
|
|
return self._chat_token_usage[key]
|
|
|
|
def _bucket_try_consume(self, chat_id: str, is_control: bool = False) -> float:
|
|
"""Consume one token. Returns 0 if available, else seconds until the next minute window.
|
|
|
|
Normal messages only use the normal quota; control messages use normal
|
|
quota first (don't waste reserved), then the reserved pool.
|
|
"""
|
|
usage = self._get_token_usage(chat_id)
|
|
now = time.monotonic()
|
|
if now - usage["last_reset"] > 60.0: # reset counters every minute
|
|
usage["normal"] = 0.0
|
|
usage["reserved"] = 0.0
|
|
usage["last_reset"] = now
|
|
if usage["normal"] < self._BUCKET_NORMAL_TOKENS:
|
|
usage["normal"] += 1.0
|
|
return 0.0
|
|
if is_control and usage["reserved"] < self._BUCKET_RESERVED_TOKENS:
|
|
usage["reserved"] += 1.0
|
|
return 0.0
|
|
return 60.0 - (now - usage["last_reset"])
|
|
|
|
async def _enqueue_chat_send(self, chat_id: str, coro_factory, is_control: bool = False):
|
|
"""Enqueue a send task for a chat and await its result (FIFO per chat, parallel across chats).
|
|
|
|
Control-lane sends bypass the normal queue so approval prompts are never blocked.
|
|
"""
|
|
key = str(chat_id or "").strip()
|
|
lane = "control" if is_control else "normal"
|
|
queues = self._control_queues if is_control else self._chat_queues
|
|
if key not in queues:
|
|
logger.debug("[%s] Creating %s queue + worker for chat %s", self.name, lane, key)
|
|
queues[key] = asyncio.Queue()
|
|
workers = self._control_workers if is_control else self._chat_workers
|
|
workers[key] = asyncio.create_task(self._send_worker(key, is_control))
|
|
queue = queues[key]
|
|
logger.debug("[%s] Enqueuing send for chat %s (lane=%s, qsize=%d)", self.name, key, lane, queue.qsize())
|
|
future = asyncio.get_running_loop().create_future()
|
|
await queue.put((coro_factory, future))
|
|
return await future
|
|
|
|
async def _send_worker(self, chat_key: str, is_control: bool) -> None:
|
|
"""Per-chat worker: drain one lane's queue under the token bucket."""
|
|
if is_control:
|
|
queue = self._control_queues[chat_key]
|
|
else:
|
|
queue = self._chat_queues[chat_key]
|
|
logger.debug("[%s] Normal send worker started for chat %s", self.name, chat_key)
|
|
try:
|
|
while True:
|
|
coro_factory, future = await queue.get()
|
|
try:
|
|
wait = self._bucket_try_consume(chat_key, is_control)
|
|
if wait > 0:
|
|
if not is_control:
|
|
logger.debug(
|
|
"[%s] Normal worker rate-limited for chat %s, waiting %.1fs",
|
|
self.name, chat_key, wait,
|
|
)
|
|
await asyncio.sleep(wait)
|
|
self._bucket_try_consume(chat_key, is_control) # re-consume after wait
|
|
result = await coro_factory()
|
|
if not future.done():
|
|
future.set_result(result)
|
|
except Exception as exc:
|
|
if not future.done():
|
|
future.set_exception(exc)
|
|
finally:
|
|
queue.task_done()
|
|
except asyncio.CancelledError:
|
|
while not queue.empty():
|
|
try:
|
|
_, future = queue.get_nowait()
|
|
if not future.done():
|
|
future.set_exception(RuntimeError("WeCom adapter shutting down"))
|
|
except asyncio.QueueEmpty:
|
|
break
|