327 lines
13 KiB
Python
327 lines
13 KiB
Python
"""AI 写作会话定期清理。
|
||
|
||
按 ``config.yaml`` 的 ``ai_writing.cleanup`` 段做联动清理：
|
||
|
||
1. 找出 MySQL ``ai_writing_sessions`` 表里 ``updated_at < now - retention_days``
|
||
的会话 id 列表（分批，避免大事务）
|
||
2. 对每个 id 调用 LangGraph checkpointer 的 ``adelete_thread()``，删掉
|
||
SQLite checkpointer 里对应的 thread 状态
|
||
3. 再把 MySQL 行批量删掉
|
||
|
||
整个流程「先删 checkpointer 再删 DB」是有意为之 —— 即便清理被中途中断，
|
||
也只会出现「checkpointer 删了但 DB 没删」的孤儿，下一轮清理会发现这些
|
||
DB 行仍然过期，重新尝试清理（checkpointer 删一个不存在的 thread 是 idempotent）。
|
||
反过来如果先删 DB，万一 checkpointer 删失败，文件里就永久留着不可达的状态。
|
||
|
||
多 worker 下用文件锁做 leader election，只有抢到锁的 worker 才会真正跑清理。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
import logging
|
||
import sys
|
||
from datetime import UTC, datetime, timedelta
|
||
from pathlib import Path
|
||
|
||
from croniter import croniter
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
# ── 文件锁 ────────────────────────────────────────────────────────────────────
|
||
|
||
|
||
def _acquire_lock(lock_path: Path):
|
||
"""跨平台非阻塞独占文件锁，多 worker 下做 leader election。
|
||
|
||
成功 → 返回持有锁的文件句柄（调用方需保持其存活直到主动释放）；
|
||
抢不到 / 打开失败 → 返回 None。进程退出时 OS 自动释放。
|
||
"""
|
||
try:
|
||
lock_path.parent.mkdir(parents=True, exist_ok=True)
|
||
handle = open(lock_path, "a+")
|
||
except OSError:
|
||
logger.exception("ai_writing_cleanup: 打开锁文件失败 path=%s", lock_path)
|
||
return None
|
||
try:
|
||
if sys.platform == "win32":
|
||
import msvcrt
|
||
|
||
handle.seek(0)
|
||
msvcrt.locking(handle.fileno(), msvcrt.LK_NBLCK, 1)
|
||
else:
|
||
import fcntl
|
||
|
||
fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
|
||
except OSError:
|
||
handle.close()
|
||
return None
|
||
return handle
|
||
|
||
|
||
def _release_lock(handle) -> None:
|
||
try:
|
||
if sys.platform == "win32":
|
||
import msvcrt
|
||
|
||
handle.seek(0)
|
||
msvcrt.locking(handle.fileno(), msvcrt.LK_UNLCK, 1)
|
||
else:
|
||
import fcntl
|
||
|
||
fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
|
||
except OSError:
|
||
pass
|
||
finally:
|
||
try:
|
||
handle.close()
|
||
except OSError:
|
||
pass
|
||
|
||
|
||
# ── 清理执行 ──────────────────────────────────────────────────────────────────
|
||
|
||
|
||
async def _delete_checkpoint_thread(checkpointer, thread_id: str) -> None:
|
||
"""删 checkpointer 里 thread_id 对应的全部状态。失败只记日志，不抛。"""
|
||
if checkpointer is None:
|
||
return
|
||
# LangGraph 异步删法首选 adelete_thread；老版本只有同步 delete_thread。
|
||
delete_async = getattr(checkpointer, "adelete_thread", None)
|
||
if delete_async is not None:
|
||
try:
|
||
await delete_async(thread_id)
|
||
return
|
||
except Exception:
|
||
logger.warning(
|
||
"ai_writing_cleanup: adelete_thread failed thread_id=%s", thread_id, exc_info=True
|
||
)
|
||
return
|
||
delete_sync = getattr(checkpointer, "delete_thread", None)
|
||
if delete_sync is not None:
|
||
try:
|
||
# 同步方法跑在线程池里，避免堵 event loop
|
||
await asyncio.to_thread(delete_sync, thread_id)
|
||
except Exception:
|
||
logger.warning(
|
||
"ai_writing_cleanup: delete_thread failed thread_id=%s", thread_id, exc_info=True
|
||
)
|
||
|
||
|
||
async def _get_active_session_ids_langgraph(thread_store) -> set[str]:
|
||
"""新路径下「正在跑 / 等用户干预」的 thread_id 集合。
|
||
|
||
``threads_meta.status`` 由 ``start_run`` / ``run_agent`` 维护：
|
||
- ``running``：graph 正在跑（节点中）
|
||
- ``interrupted``：graph 暂停在 interrupt 点等用户回头填卡
|
||
这两种都视作活跃，**绝不能删**。``idle / success / error`` 才是真正终止状态。
|
||
|
||
LangGraph Server 路径 + 老路径的 lead_agent 都用同一份 ``thread_store``，所以
|
||
这里拿的 active 集合是全部 graph 的；筛 ai_writing 是按 assistant_id 过滤。
|
||
"""
|
||
if thread_store is None:
|
||
return set()
|
||
try:
|
||
running = await thread_store.search(status="running", user_id=None, limit=1000)
|
||
interrupted = await thread_store.search(status="interrupted", user_id=None, limit=1000)
|
||
except Exception:
|
||
logger.warning("ai_writing_cleanup: 读取 LangGraph 活动 thread 列表失败", exc_info=True)
|
||
return set()
|
||
# 仅关心 assistant_id='ai_writing' 的 thread —— 别的 graph 不归这套清理管
|
||
active: set[str] = set()
|
||
for row in running + interrupted:
|
||
if row.get("assistant_id") == "ai_writing":
|
||
active.add(row["thread_id"])
|
||
return active
|
||
|
||
|
||
async def run_cleanup_once(
|
||
*,
|
||
retention_days: int,
|
||
only_finished: bool = False,
|
||
batch_size: int = 500,
|
||
checkpointer=None,
|
||
thread_store=None,
|
||
) -> dict:
|
||
"""执行一轮清理。返回统计信息。
|
||
|
||
Args:
|
||
checkpointer: 主 LangGraph checkpointer（``app.state.checkpointer``）。
|
||
AI 写作的 thread 状态由 LangGraph Server 注入，存在这里 —— admin 路由
|
||
+ scheduler loop 都会从 ``app.state`` 取到引用传进来；测试场景下传 None
|
||
会让 checkpoint 删除变成 no-op，但 DB 行仍会被清。
|
||
thread_store: ``app.state.thread_store``。用来按 ``status in (running,
|
||
interrupted)`` 查活跃 thread 列表做活跃保护。
|
||
|
||
**安全保证**：清理永远跳过仍在跑 / 等用户干预的会话（``threads_meta.status
|
||
in (running, interrupted)`` + ``assistant_id='ai_writing'``）。即使 admin
|
||
手动触发 ``retention_days=0`` 也强制生效。
|
||
|
||
幂等 + 可重入 —— 多次调用对结果只会越扫越干净。即便清理途中被中断，下次
|
||
再跑会重新把 DB 里仍然过期的行扫到（checkpointer 删一个已不存在的 thread
|
||
是 no-op）。
|
||
"""
|
||
from deerflow.persistence.ai_writing_sessions import AIWritingSessionRepository
|
||
from deerflow.persistence.engine import get_session_factory
|
||
|
||
sf = get_session_factory()
|
||
if sf is None:
|
||
logger.info("ai_writing_cleanup: persistence backend=memory，跳过清理")
|
||
return {"deleted_db": 0, "deleted_checkpoints": 0, "skipped_active": 0}
|
||
|
||
repo = AIWritingSessionRepository(sf)
|
||
active_ids = await _get_active_session_ids_langgraph(thread_store)
|
||
|
||
cutoff = datetime.now(UTC) - timedelta(days=retention_days)
|
||
logger.info(
|
||
"ai_writing_cleanup: 开始清理 cutoff=%s only_finished=%s batch_size=%d active=%d",
|
||
cutoff.isoformat(),
|
||
only_finished,
|
||
batch_size,
|
||
len(active_ids),
|
||
)
|
||
|
||
total_db = 0
|
||
total_cp = 0
|
||
total_skipped = 0
|
||
while True:
|
||
ids = await repo.list_older_than(cutoff, only_finished=only_finished, limit=batch_size)
|
||
if not ids:
|
||
break
|
||
|
||
# 关键保护：跳过任何仍在内存里跑的会话
|
||
ids_to_delete = [sid for sid in ids if sid not in active_ids]
|
||
skipped = len(ids) - len(ids_to_delete)
|
||
if skipped:
|
||
logger.warning(
|
||
"ai_writing_cleanup: 跳过 %d 个仍在跑的会话: %s",
|
||
skipped,
|
||
[sid for sid in ids if sid in active_ids],
|
||
)
|
||
total_skipped += skipped
|
||
|
||
if not ids_to_delete:
|
||
# 本批全是活动会话 —— 但 list_older_than 是按 updated_at 升序的，
|
||
# 下一批拿出来的更旧，所以这种情况只可能整个 cutoff 区间都被活动
|
||
# 会话占满，不用继续往下扫，break
|
||
break
|
||
|
||
# 先删 checkpointer，再删 DB —— 保证「DB 删了 → checkpointer 一定也清干净」
|
||
for sid in ids_to_delete:
|
||
await _delete_checkpoint_thread(checkpointer, sid)
|
||
total_cp += 1
|
||
|
||
deleted = await repo.delete_by_ids(ids_to_delete)
|
||
total_db += deleted
|
||
logger.info(
|
||
"ai_writing_cleanup: 本批删 %d 行（DB），%d 个 thread（checkpointer），跳过 %d",
|
||
deleted,
|
||
len(ids_to_delete),
|
||
skipped,
|
||
)
|
||
|
||
# 整批比 batch_size 小 → 没有更多了
|
||
if len(ids) < batch_size:
|
||
break
|
||
|
||
logger.info(
|
||
"ai_writing_cleanup: 清理完成 total_db=%d total_checkpoints=%d skipped_active=%d",
|
||
total_db,
|
||
total_cp,
|
||
total_skipped,
|
||
)
|
||
return {
|
||
"deleted_db": total_db,
|
||
"deleted_checkpoints": total_cp,
|
||
"skipped_active": total_skipped,
|
||
}
|
||
|
||
|
||
# ── 调度循环 ──────────────────────────────────────────────────────────────────
|
||
|
||
|
||
def _next_cron_delay_seconds(cron_expr: str, base: datetime | None = None) -> float:
|
||
"""根据 cron 表达式算到下一次触发的秒数。
|
||
|
||
用 ``datetime.now(UTC)`` 做基准。croniter 接受 naive 时间，这里转成 naive
|
||
UTC 喂进去；当然如果用户希望按本地时区跑（凌晨 3 点），需要在外层调整 ——
|
||
我们用本地时区基准，等价于 cron 在本地时区跑。
|
||
"""
|
||
base = base or datetime.now()
|
||
nxt = croniter(cron_expr, base).get_next(datetime)
|
||
return max(1.0, (nxt - base).total_seconds())
|
||
|
||
|
||
async def cleanup_scheduler_loop(app=None) -> None:
|
||
"""常驻后台任务：按 cron 调度跑清理。
|
||
|
||
多 worker 部署下文件锁做 leader election，没拿到锁的 worker 直接退出。
|
||
被 cancel 时 CancelledError 不会被 except Exception 吞掉，能干净退出。
|
||
|
||
Args:
|
||
app: FastAPI 实例，用来从 ``app.state`` 拿主 checkpointer / thread_store
|
||
引用，喂给 ``run_cleanup_once`` —— 新路径下 ai_writing 的 thread 状态
|
||
在主 checkpointer 里。``app=None`` 时回退到老 checkpointer（仅老路径）。
|
||
"""
|
||
from deerflow.config.app_config import get_app_config
|
||
from deerflow.config.runtime_paths import runtime_home
|
||
|
||
try:
|
||
cfg = get_app_config()
|
||
except Exception:
|
||
logger.exception("ai_writing_cleanup: 加载 app config 失败，调度器退出")
|
||
return
|
||
|
||
raw = getattr(cfg, "ai_writing", None) or {}
|
||
cleanup_cfg = (raw.get("cleanup") if isinstance(raw, dict) else None) or {}
|
||
if not cleanup_cfg.get("enabled", False):
|
||
logger.info("ai_writing_cleanup: cleanup.enabled=false，调度器不启动")
|
||
return
|
||
|
||
cron_expr = str(cleanup_cfg.get("cron", "0 3 * * *"))
|
||
retention_days = int(cleanup_cfg.get("retention_days", 7))
|
||
only_finished = bool(cleanup_cfg.get("delete_only_finished", False))
|
||
batch_size = int(cleanup_cfg.get("batch_size", 500))
|
||
|
||
# 提前校验 cron 表达式，配错就早 fail 早暴露
|
||
try:
|
||
croniter(cron_expr, datetime.now())
|
||
except Exception:
|
||
logger.exception("ai_writing_cleanup: 非法 cron 表达式 %r，调度器退出", cron_expr)
|
||
return
|
||
|
||
lock_handle = _acquire_lock(runtime_home() / ".ai_writing_cleanup.lock")
|
||
if lock_handle is None:
|
||
logger.info("ai_writing_cleanup: 另一个 worker 已持有调度锁，本 worker 跳过")
|
||
return
|
||
|
||
logger.info(
|
||
"ai_writing_cleanup: 已启动 cron=%r retention_days=%d only_finished=%s batch_size=%d",
|
||
cron_expr,
|
||
retention_days,
|
||
only_finished,
|
||
batch_size,
|
||
)
|
||
try:
|
||
while True:
|
||
delay = _next_cron_delay_seconds(cron_expr)
|
||
logger.info("ai_writing_cleanup: 下一次清理 %.0f 秒后", delay)
|
||
await asyncio.sleep(delay)
|
||
try:
|
||
# 每次循环重新读 app.state —— 不缓存，避免 lifespan 中重启 / 重置
|
||
# checkpointer 时拿到死引用。app=None 时 main_checkpointer / store
|
||
# 也为 None，run_cleanup_once 内部走老 fallback。
|
||
main_checkpointer = getattr(getattr(app, "state", None), "checkpointer", None)
|
||
main_thread_store = getattr(getattr(app, "state", None), "thread_store", None)
|
||
await run_cleanup_once(
|
||
retention_days=retention_days,
|
||
only_finished=only_finished,
|
||
batch_size=batch_size,
|
||
checkpointer=main_checkpointer,
|
||
thread_store=main_thread_store,
|
||
)
|
||
except Exception:
|
||
logger.exception("ai_writing_cleanup: 本轮清理失败（不影响下一轮）")
|
||
finally:
|
||
_release_lock(lock_handle)
|