vent_agent_2.0/scheduler/runner.py

"""定时任务执行器。

一次执行 = 建一个独立会话 + 跑一轮 agent + 把结果落库并计费。

为什么用「独立会话」而不是 1.5 那种 ``scheduled-task-{id}-{run_id}`` 裸线程 ID：
本系统其他所有地方都强制 ``thread_id ≡ session_id``（那条不变式是防止
跨会话读写图状态的第一道闸）。定时任务若用另一个 ID 命名法，就成了唯一
的例外，将来谁改了都会踩。用独立会话既同样隔离，又让用户能在会话列表里
直接翻到定时任务的产出——1.5 的定时任务结果只能去 schedule_runs 表里看。
"""

from __future__ import annotations

import asyncio
from typing import Any

import config
from billing import pricing
from db import chat_store, db_conn, schedule_store
from obs import get_logger

logger = get_logger(__name__)

# 定时任务的图递归上限：比交互式对话紧得多。
# 理由是没人盯着它——跑飞了只会白烧额度，早点中止更划算。
SCHEDULED_RECURSION_LIMIT = 30


async def run_schedule(schedule_id: str, trigger: str = "cron") -> dict[str, Any]:
    """执行并持久化状态；准入失败由 HTTP 或调度回调处理。"""
    schedule = await asyncio.to_thread(schedule_store.get_schedule, schedule_id)
    if schedule is None:
        logger.warning("定时任务不存在 id=%s", schedule_id)
        return {"ok": False, "error": "schedule_not_found"}

    from agents import execution, prompts
    from api import concurrency
    from contextlib import aclosing
    import anyio

    user_name = schedule.get("user_name") or "default"
    # 同一个计划的 cron 和手动执行共用互斥，避免重复触发。
    with concurrency.session_guard(f"schedule:{schedule_id}"):
        from billing import guard
        await asyncio.to_thread(guard.precheck, user_name)
        session = await asyncio.to_thread(chat_store.create_session, user_name,
            schedule.get("mode") or "full", f"[定时] {schedule.get('title') or schedule_id[:8]}")
        thread_id = session["id"]
        admission = await execution.admit(user_name, thread_id)
        run_id = None
        stream = None
        try:
            run_id = await asyncio.to_thread(schedule_store.start_run, schedule_id, thread_id)
            prompt = schedule.get("prompt") or ""
            preferences = await asyncio.to_thread(chat_store.preferences_as_text, user_name)
            message = prompt + prompts.build_preference_instruction(preferences) + prompts.build_mode_instruction(schedule.get("mode") or "full")
            stream = await execution.build_stream(admission,
                session_id=thread_id, thread_id=thread_id, user_name=user_name,
                kind="scheduled", model=config.MODEL, original_user_message=prompt, message=message,
                config={"configurable": {"thread_id": thread_id}, "recursion_limit": SCHEDULED_RECURSION_LIMIT})
            async with aclosing(stream.generate()) as events:
                async for _ in events:
                    pass
            output = "".join(stream.answer_parts)
            await asyncio.to_thread(schedule_store.finish_run, run_id, stream.status,
                                    stream.error or None, output[:2000])
            return {"ok": stream.status in {"success", "waiting_input"}, "status": stream.status,
                    "run_id": run_id, "session_id": thread_id, "output": output,
                    "error": stream.error, "input_tokens": stream.meter.input_tokens,
                    "output_tokens": stream.meter.output_tokens}
        except BaseException as exc:
            if run_id:
                with anyio.CancelScope(shield=True):
                    await asyncio.to_thread(schedule_store.finish_run, run_id,
                        "cancelled" if isinstance(exc, asyncio.CancelledError) else "failed", type(exc).__name__, None)
            raise
        finally:
            admission.release()


def ensure_pricing_seeded() -> None:
    """确保单价表就绪，避免定时任务因为模型没定价而 fail-closed。"""
    try:
        pricing.seed_default_pricing()
    except Exception as exc:  # pragma: no cover - 取决于数据库可用性
        logger.warning("单价表初始化失败：%s: %s", type(exc).__name__, exc)


__all__ = ["SCHEDULED_RECURSION_LIMIT", "ensure_pricing_seeded", "run_schedule", "db_conn"]