66 lines
1.9 KiB
Python
66 lines
1.9 KiB
Python
import asyncio
|
|
from typing import Optional, Callable
|
|
from datetime import datetime, timedelta
|
|
|
|
from .config import Config
|
|
|
|
|
|
class MessageService:
|
|
def __init__(self, config: Config):
|
|
self.config = config
|
|
self.pending_recalls: dict[str, list[asyncio.Task]] = {}
|
|
|
|
async def send_with_recall(
|
|
self,
|
|
scope: str,
|
|
message_type: str,
|
|
send_func: Callable,
|
|
*args,
|
|
**kwargs,
|
|
) -> Optional[str]:
|
|
"""Send message and schedule recall if configured."""
|
|
try:
|
|
message_id = await send_func(*args, **kwargs)
|
|
if not message_id:
|
|
return None
|
|
|
|
recall_time = self.config.MESSAGE_RECALL.get(message_type, 0)
|
|
if recall_time > 0:
|
|
task = asyncio.create_task(
|
|
self._schedule_recall(scope, message_id, recall_time, send_func)
|
|
)
|
|
if scope not in self.pending_recalls:
|
|
self.pending_recalls[scope] = []
|
|
self.pending_recalls[scope].append(task)
|
|
|
|
return message_id
|
|
except Exception as e:
|
|
return None
|
|
|
|
async def _schedule_recall(
|
|
self,
|
|
scope: str,
|
|
message_id: str,
|
|
delay: int,
|
|
recall_func: Callable,
|
|
):
|
|
"""Schedule message recall."""
|
|
try:
|
|
await asyncio.sleep(delay)
|
|
await recall_func(message_id)
|
|
except Exception:
|
|
pass
|
|
|
|
def clear_pending_recalls(self, scope: str):
|
|
"""Cancel all pending recall tasks for a scope."""
|
|
if scope in self.pending_recalls:
|
|
for task in self.pending_recalls[scope]:
|
|
if not task.done():
|
|
task.cancel()
|
|
del self.pending_recalls[scope]
|
|
|
|
def clear_all_recalls(self):
|
|
"""Cancel all pending recall tasks."""
|
|
for scope in list(self.pending_recalls.keys()):
|
|
self.clear_pending_recalls(scope)
|