RRRosmontis
34089d0042
Include utility modules, shared channels, prompts, read-only skills, wardrobe memory and regression tests. Externalize legacy credentials and exclude private runtime state.
406 lines
17 KiB
Python
406 lines
17 KiB
Python
"""Persistent once, daily, and weekly reminders backed by SQLite."""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import json
|
||
import os
|
||
import re
|
||
import secrets
|
||
import sqlite3
|
||
import threading
|
||
import time
|
||
from datetime import datetime, timedelta
|
||
from pathlib import Path
|
||
from typing import Any, Mapping
|
||
from zoneinfo import ZoneInfo
|
||
|
||
import notifications
|
||
|
||
SHANGHAI = ZoneInfo("Asia/Shanghai")
|
||
BASE_DIR = Path(__file__).resolve().parent
|
||
DB_PATH = os.environ.get("AGENT_DB_PATH", str(BASE_DIR / "data" / "agent.db"))
|
||
_ALLOWED_PARAMS = {"message", "kind", "run_at", "time_of_day", "weekdays"}
|
||
_TIME_RE = re.compile(r"^(\d{2}):(\d{2})(?::(\d{2}))?$")
|
||
_KIND_NAMES = {"once": "单次", "daily": "每天", "weekly": "每周"}
|
||
FIRING_LEASE_SECONDS = 300.0
|
||
|
||
|
||
def _epoch(value: datetime | float | int | None) -> float:
|
||
if value is None:
|
||
return time.time()
|
||
if isinstance(value, datetime):
|
||
if value.tzinfo is None:
|
||
value = value.replace(tzinfo=SHANGHAI)
|
||
return value.timestamp()
|
||
if isinstance(value, bool) or not isinstance(value, (int, float)):
|
||
raise TypeError("now 必须是 datetime 或 epoch 秒数")
|
||
return float(value)
|
||
|
||
|
||
def _local_datetime(value: float) -> datetime:
|
||
return datetime.fromtimestamp(value, SHANGHAI)
|
||
|
||
|
||
def _time_parts(value: Any) -> tuple[int, int, int]:
|
||
if not isinstance(value, str):
|
||
raise ValueError("time_of_day 必须是 HH:MM 或 HH:MM:SS")
|
||
match = _TIME_RE.fullmatch(value.strip())
|
||
if not match:
|
||
raise ValueError("time_of_day 必须是 HH:MM 或 HH:MM:SS")
|
||
hour, minute = int(match[1]), int(match[2])
|
||
second = int(match[3] or 0)
|
||
if hour > 23 or minute > 59 or second > 59:
|
||
raise ValueError("time_of_day 不是有效时间")
|
||
return hour, minute, second
|
||
|
||
|
||
def _next_daily(now: float, time_of_day: Any) -> float:
|
||
hour, minute, second = _time_parts(time_of_day)
|
||
local = _local_datetime(now)
|
||
candidate = local.replace(hour=hour, minute=minute, second=second, microsecond=0)
|
||
if candidate.timestamp() <= now:
|
||
candidate += timedelta(days=1)
|
||
return candidate.timestamp()
|
||
|
||
|
||
def _weekdays(value: Any) -> list[int]:
|
||
if not isinstance(value, list) or not value:
|
||
raise ValueError("weekly 必须提供非空 weekdays 列表")
|
||
if any(isinstance(day, bool) or not isinstance(day, int) or not 1 <= day <= 7 for day in value):
|
||
raise ValueError("weekdays 只能包含 1..7 的整数")
|
||
return sorted(set(value))
|
||
|
||
|
||
def _next_weekly(now: float, time_of_day: Any, weekdays: Any) -> float:
|
||
days = _weekdays(weekdays)
|
||
hour, minute, second = _time_parts(time_of_day)
|
||
local = _local_datetime(now)
|
||
midnight = local.replace(hour=0, minute=0, second=0, microsecond=0)
|
||
for offset in range(8):
|
||
day = midnight + timedelta(days=offset)
|
||
if day.isoweekday() not in days:
|
||
continue
|
||
candidate = day.replace(hour=hour, minute=minute, second=second)
|
||
if candidate.timestamp() > now:
|
||
return candidate.timestamp()
|
||
raise AssertionError("weekly next run calculation failed")
|
||
|
||
|
||
def _next_run(kind: str, now: float, params: Mapping[str, Any]) -> float:
|
||
if kind == "daily":
|
||
return _next_daily(now, params.get("time_of_day"))
|
||
if kind == "weekly":
|
||
return _next_weekly(now, params.get("time_of_day"), params.get("weekdays"))
|
||
raise ValueError("不支持的重复类型")
|
||
|
||
|
||
def _epoch_text(value: float) -> str:
|
||
if float(value).is_integer():
|
||
return str(int(value))
|
||
return f"{value:.6f}".rstrip("0").rstrip(".")
|
||
|
||
|
||
def _row_dict(row: sqlite3.Row | Mapping[str, Any]) -> dict[str, Any]:
|
||
task = dict(row)
|
||
if task.get("weekdays"):
|
||
task["weekdays"] = json.loads(task["weekdays"])
|
||
else:
|
||
task["weekdays"] = None
|
||
next_run_at = task.get("next_run_at")
|
||
task["next_run_text"] = (
|
||
_local_datetime(float(next_run_at)).strftime("%Y-%m-%d %H:%M:%S")
|
||
if next_run_at is not None else None
|
||
)
|
||
return task
|
||
|
||
|
||
def format_task(task: Mapping[str, Any]) -> str:
|
||
"""Render one task as concise Chinese text for a user."""
|
||
kind = str(task.get("kind", ""))
|
||
next_run_at = task.get("next_run_at")
|
||
next_text = task.get("next_run_text")
|
||
if next_text is None and next_run_at is not None:
|
||
next_text = _local_datetime(float(next_run_at)).strftime("%Y-%m-%d %H:%M:%S")
|
||
return (
|
||
f"ID：{task.get('id', '')}\n"
|
||
f"类型：{_KIND_NAMES.get(kind, kind)}\n"
|
||
f"消息：{task.get('message', '')}\n"
|
||
f"下次时间：{next_text or '无'}"
|
||
)
|
||
|
||
|
||
HELP_TEXT = """定时提醒用法：
|
||
- 创建：直接用自然语言说“今天下午6点提醒我开会”（单次）、“每天08:30提醒我签到”（每天）或“每周一、周五18:00提醒我写周报”（每周）。也可以提供 message、kind、run_at、time_of_day、weekdays 参数。
|
||
- 查看：说“查看我的定时提醒”或“提醒列表”，只显示当前私聊或群聊中的有效提醒。
|
||
- 取消：说“取消提醒 <ID>”，ID 可以只写能唯一匹配当前会话任务的前缀。
|
||
时间均按 Asia/Shanghai（北京时间）计算。"""
|
||
|
||
|
||
class Scheduler:
|
||
def __init__(self, db_path: str = DB_PATH):
|
||
self.db_path = str(db_path)
|
||
self._lock = threading.RLock()
|
||
parent = os.path.dirname(os.path.abspath(self.db_path))
|
||
os.makedirs(parent, exist_ok=True)
|
||
self._init_db()
|
||
|
||
def _connect(self) -> sqlite3.Connection:
|
||
con = sqlite3.connect(self.db_path, timeout=30)
|
||
con.row_factory = sqlite3.Row
|
||
con.execute("PRAGMA journal_mode=WAL")
|
||
return con
|
||
|
||
def _init_db(self) -> None:
|
||
with self._lock, self._connect() as con:
|
||
con.execute(
|
||
"""CREATE TABLE IF NOT EXISTS scheduled_tasks(
|
||
id TEXT PRIMARY KEY,
|
||
uid INTEGER,
|
||
channel TEXT,
|
||
msg_type TEXT,
|
||
target_id INTEGER,
|
||
message TEXT NOT NULL,
|
||
kind TEXT NOT NULL,
|
||
run_at REAL,
|
||
time_of_day TEXT,
|
||
weekdays TEXT,
|
||
next_run_at REAL,
|
||
created_at REAL NOT NULL,
|
||
updated_at REAL NOT NULL,
|
||
status TEXT NOT NULL DEFAULT 'active',
|
||
firing_owner TEXT,
|
||
firing_until REAL
|
||
)"""
|
||
)
|
||
task_columns = {
|
||
row["name"] for row in con.execute("PRAGMA table_info(scheduled_tasks)")
|
||
}
|
||
for name, declaration in (("firing_owner", "TEXT"), ("firing_until", "REAL")):
|
||
if name not in task_columns:
|
||
con.execute(f"ALTER TABLE scheduled_tasks ADD COLUMN {name} {declaration}")
|
||
con.execute(
|
||
"""CREATE INDEX IF NOT EXISTS scheduled_tasks_due
|
||
ON scheduled_tasks(status,next_run_at)"""
|
||
)
|
||
|
||
@staticmethod
|
||
def _context(ctx: Any) -> tuple[str, int | None, str, str, int]:
|
||
channel = str(ctx.channel.name)
|
||
msg_type = str(ctx.msg_type)
|
||
target_id = int(ctx.target_id)
|
||
uid = int(ctx.user_id) if msg_type == "private" else None
|
||
return ("private" if uid is not None else "group", uid, channel, msg_type, target_id)
|
||
|
||
@staticmethod
|
||
def _scope(ctx: Any) -> tuple[str, tuple[Any, ...]]:
|
||
if ctx.msg_type == "private":
|
||
return "uid=?", (int(ctx.user_id),)
|
||
return "channel=? AND msg_type=? AND target_id=?", (
|
||
str(ctx.channel.name), str(ctx.msg_type), int(ctx.target_id)
|
||
)
|
||
|
||
def create_from_params(
|
||
self, ctx: Any, params: Mapping[str, Any], now: datetime | float | int | None = None
|
||
) -> dict[str, Any]:
|
||
if not isinstance(params, Mapping):
|
||
raise ValueError("params 必须是对象")
|
||
unknown = set(params) - _ALLOWED_PARAMS
|
||
if unknown:
|
||
raise ValueError(f"不支持的参数：{', '.join(sorted(map(str, unknown)))}")
|
||
message = params.get("message")
|
||
if not isinstance(message, str) or not message.strip() or len(message) > 1000:
|
||
raise ValueError("message 必须是非空且不超过 1000 字的字符串")
|
||
message = message.strip()
|
||
kind = params.get("kind")
|
||
if kind not in ("once", "daily", "weekly"):
|
||
raise ValueError("kind 必须是 once、daily 或 weekly")
|
||
current = _epoch(now)
|
||
run_at = None
|
||
time_of_day = None
|
||
weekdays = None
|
||
if kind == "once":
|
||
raw_run_at = params.get("run_at")
|
||
if not isinstance(raw_run_at, str) or not raw_run_at.strip():
|
||
raise ValueError("once 必须提供 ISO8601 run_at")
|
||
try:
|
||
parsed = datetime.fromisoformat(raw_run_at.strip())
|
||
except ValueError as exc:
|
||
raise ValueError("run_at 必须是有效的 ISO8601 时间") from exc
|
||
if parsed.tzinfo is None:
|
||
parsed = parsed.replace(tzinfo=SHANGHAI)
|
||
run_at = parsed.timestamp()
|
||
if run_at <= current:
|
||
raise ValueError("run_at 必须是未来时间")
|
||
next_run_at = run_at
|
||
else:
|
||
time_of_day = params.get("time_of_day")
|
||
# Validate and canonicalize HH:MM to HH:MM:SS for persistent storage.
|
||
hour, minute, second = _time_parts(time_of_day)
|
||
time_of_day = f"{hour:02d}:{minute:02d}:{second:02d}"
|
||
if kind == "daily":
|
||
next_run_at = _next_daily(current, time_of_day)
|
||
else:
|
||
days = _weekdays(params.get("weekdays"))
|
||
weekdays = json.dumps(days)
|
||
next_run_at = _next_weekly(current, time_of_day, days)
|
||
|
||
_, uid, channel, msg_type, target_id = self._context(ctx)
|
||
with self._lock, self._connect() as con:
|
||
while True:
|
||
task_id = secrets.token_hex(4)
|
||
if not con.execute("SELECT 1 FROM scheduled_tasks WHERE id=?", (task_id,)).fetchone():
|
||
break
|
||
con.execute(
|
||
"""INSERT INTO scheduled_tasks(
|
||
id,uid,channel,msg_type,target_id,message,kind,run_at,
|
||
time_of_day,weekdays,next_run_at,created_at,updated_at,status
|
||
) VALUES (?,?,?,?,?,?,?,?,?,?,?,?,?,'active')""",
|
||
(task_id, uid, channel, msg_type, target_id, message, kind, run_at,
|
||
time_of_day, weekdays, next_run_at, current, current),
|
||
)
|
||
row = con.execute("SELECT * FROM scheduled_tasks WHERE id=?", (task_id,)).fetchone()
|
||
return _row_dict(row)
|
||
|
||
def list_for_context(self, ctx: Any) -> list[dict[str, Any]]:
|
||
where, args = self._scope(ctx)
|
||
with self._lock, self._connect() as con:
|
||
rows = con.execute(
|
||
f"""SELECT * FROM scheduled_tasks
|
||
WHERE status='active' AND {where}
|
||
ORDER BY next_run_at,id""",
|
||
args,
|
||
).fetchall()
|
||
return [_row_dict(row) for row in rows]
|
||
|
||
def cancel_for_context(self, ctx: Any, id_prefix: str) -> dict[str, Any]:
|
||
if not isinstance(id_prefix, str) or not id_prefix.strip():
|
||
raise ValueError("必须提供任务 ID 前缀")
|
||
prefix = id_prefix.strip()
|
||
where, args = self._scope(ctx)
|
||
with self._lock, self._connect() as con:
|
||
rows = con.execute(
|
||
f"SELECT * FROM scheduled_tasks WHERE status='active' AND {where}", args
|
||
).fetchall()
|
||
matches = [row for row in rows if str(row["id"]).startswith(prefix)]
|
||
if len(matches) != 1:
|
||
raise ValueError("任务不存在、无权取消或 ID 前缀不唯一")
|
||
task_id = matches[0]["id"]
|
||
cursor = con.execute(
|
||
"""UPDATE scheduled_tasks SET status='cancelled',updated_at=?
|
||
WHERE id=? AND status='active'""",
|
||
(time.time(), task_id),
|
||
)
|
||
if cursor.rowcount != 1:
|
||
raise ValueError("任务已不可取消")
|
||
row = con.execute("SELECT * FROM scheduled_tasks WHERE id=?", (task_id,)).fetchone()
|
||
return _row_dict(row)
|
||
|
||
def _claim_task(self, row: sqlite3.Row, current: float, owner: str) -> bool:
|
||
"""Serialize firing with cancellation and recover leases left by crashes."""
|
||
with self._lock, self._connect() as con:
|
||
cursor = con.execute(
|
||
"""UPDATE scheduled_tasks
|
||
SET status='firing',firing_owner=?,firing_until=?,updated_at=?
|
||
WHERE id=? AND next_run_at=?
|
||
AND (status='active'
|
||
OR (status='firing' AND firing_until<=?))""",
|
||
(owner, current + FIRING_LEASE_SECONDS, current,
|
||
row["id"], row["next_run_at"], current),
|
||
)
|
||
return cursor.rowcount == 1
|
||
|
||
def _release_task(self, task_id: str, owner: str, current: float) -> None:
|
||
with self._lock, self._connect() as con:
|
||
con.execute(
|
||
"""UPDATE scheduled_tasks
|
||
SET status='active',firing_owner=NULL,firing_until=NULL,updated_at=?
|
||
WHERE id=? AND status='firing' AND firing_owner=?""",
|
||
(current, task_id, owner),
|
||
)
|
||
|
||
def fire_due(
|
||
self,
|
||
now: datetime | float | int | None = None,
|
||
limit: int = 20,
|
||
center: Any = None,
|
||
) -> int:
|
||
current = _epoch(now)
|
||
if isinstance(limit, bool) or not isinstance(limit, int) or limit < 1:
|
||
raise ValueError("limit 必须是正整数")
|
||
publisher = notifications.center if center is None else center
|
||
with self._lock, self._connect() as con:
|
||
rows = con.execute(
|
||
"""SELECT * FROM scheduled_tasks
|
||
WHERE next_run_at<=?
|
||
AND (status='active'
|
||
OR (status='firing' AND firing_until<=?))
|
||
ORDER BY next_run_at,id LIMIT ?""",
|
||
(current, current, limit),
|
||
).fetchall()
|
||
fired = 0
|
||
for row in rows:
|
||
scheduled_at = float(row["next_run_at"])
|
||
owner = secrets.token_hex(16)
|
||
# A cancellation that commits after SELECT but before this atomic
|
||
# claim wins; such a task must never reach publish().
|
||
if not self._claim_task(row, current, owner):
|
||
continue
|
||
try:
|
||
if row["kind"] == "once":
|
||
next_run_at = scheduled_at
|
||
next_status = "completed"
|
||
else:
|
||
next_run_at = _next_run(
|
||
row["kind"], current,
|
||
{"time_of_day": row["time_of_day"],
|
||
"weekdays": json.loads(row["weekdays"] or "[]")},
|
||
)
|
||
next_status = "active"
|
||
event = notifications.Event(
|
||
kind="reminder",
|
||
dedupe_key=f"reminder:{row['id']}:{_epoch_text(scheduled_at)}",
|
||
uid=row["uid"] if row["msg_type"] == "private" else None,
|
||
channel=row["channel"],
|
||
msg_type=row["msg_type"],
|
||
target_id=row["target_id"],
|
||
legacy_text="定时提醒：" + row["message"],
|
||
payload={
|
||
"task_id": row["id"],
|
||
"message": row["message"],
|
||
"kind": row["kind"],
|
||
"scheduled_at": int(scheduled_at) if scheduled_at.is_integer() else scheduled_at,
|
||
},
|
||
)
|
||
event_id = publisher.publish(event)
|
||
except Exception:
|
||
self._release_task(row["id"], owner, current)
|
||
raise
|
||
if not event_id:
|
||
# No enabled target means the producer did not accept work.
|
||
# Keep the occurrence due so a later binding can accept it.
|
||
self._release_task(row["id"], owner, current)
|
||
continue
|
||
with self._lock, self._connect() as con:
|
||
cursor = con.execute(
|
||
"""UPDATE scheduled_tasks
|
||
SET next_run_at=?,status=?,updated_at=?,
|
||
firing_owner=NULL,firing_until=NULL
|
||
WHERE id=? AND status='firing' AND firing_owner=?
|
||
AND next_run_at=?""",
|
||
(next_run_at, next_status, current, row["id"], owner, scheduled_at),
|
||
)
|
||
if cursor.rowcount == 1:
|
||
fired += 1
|
||
return fired
|
||
|
||
async def worker(self, interval: float = 2) -> None:
|
||
while True:
|
||
try:
|
||
self.fire_due()
|
||
except Exception as exc:
|
||
print(f"[scheduler] worker: {type(exc).__name__}: {exc}")
|
||
await asyncio.sleep(interval)
|
||
|
||
|
||
scheduler = Scheduler()
|