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
import logging
import time
from typing import Any

import config
from billing import pricing, guard
from runtime import concurrency
from db import chat_store, db_conn, schedule_store
from obs import get_logger, get_trace_id, bind_trace_id, unbind_trace_id, log_context

logger = get_logger(__name__)

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


async def run_schedule(schedule_id: str, trigger: str = "cron") -> dict[str, Any]:
    token = bind_trace_id(get_trace_id())
    started = time.perf_counter()
    try:
        with log_context(schedule_id=schedule_id, trigger=trigger):
            try:
                result = await _run_schedule(schedule_id, trigger)
            except asyncio.CancelledError:
                logger.warning("定时执行已取消", extra={"kv": {"event": "schedule.cancelled"}})
                raise
            except (guard.QuotaExceededError, concurrency.TooManyTasksError, concurrency.SessionBusyError) as exc:
                logger.warning("定时执行未获准", extra={"kv": {"event": "schedule.rejected", "error_type": type(exc).__name__}})
                raise
            except Exception:
                logger.exception("定时执行失败", extra={"kv": {"event": "schedule.failed"}})
                raise
            logger.log(logging.INFO if result.get("ok") else logging.WARNING, "定时执行结束", extra={"kv": {
                "event": "schedule.completed", "status": result.get("status", result.get("error")),
                "schedule_run_id": result.get("run_id"), "session_id": result.get("session_id"),
                "duration_ms": round((time.perf_counter() - started) * 1000, 1),
            }})
            return result
    finally:
        unbind_trace_id(token)


async def _run_schedule(schedule_id: str, trigger: str) -> 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 prompts
    from runtime import execution
    from runtime import concurrency
    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_execution(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})
            result = await stream.run()
            await asyncio.to_thread(schedule_store.finish_run, run_id, result.status,
                                    result.error or None, result.output[:2000])
            return {"ok": result.status in {"success", "waiting_input"}, "status": result.status,
                    "run_id": run_id, "session_id": thread_id, "output": result.output,
                    "error": result.error, "input_tokens": result.usage["input_tokens"],
                    "output_tokens": result.usage["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"]