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

from langchain_core.messages import AIMessage, HumanMessage

import config
from api.context_meter import extract_usage
from billing import pricing
from billing.meter import UsageMeter
from db import chat_store, db_conn, schedule_store
from obs import get_logger

logger = get_logger(__name__)

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


def _extract_final_text(result: dict[str, Any]) -> str:
    """从图输出里取最后一条 AI 回复的正文。"""
    for message in reversed(result.get("messages") or []):
        if isinstance(message, AIMessage):
            content = message.content
            if isinstance(content, str) and content.strip():
                return content
            if isinstance(content, list):
                text = "".join(
                    block.get("text", "")
                    for block in content
                    if isinstance(block, dict)
                )
                if text.strip():
                    return text
    return ""


def _sum_usage(messages: list[Any]) -> tuple[int, int, str]:
    """汇总本次执行的真实 token 用量。"""
    input_tokens = output_tokens = 0
    source = "provider"
    saw_provider = False
    for message in messages or []:
        usage = extract_usage(message)
        if usage is None:
            continue
        saw_provider = True
        input_tokens += usage["input_tokens"]
        output_tokens += usage["output_tokens"]
    if not saw_provider and messages:
        # provider 没给用量：标 estimate 而不是静默按 0 记账
        source = "estimate"
    return input_tokens, output_tokens, source


async def run_schedule(schedule_id: str, trigger: str = "cron") -> dict[str, Any]:
    """执行一次定时任务。异常不外抛——调度器的回调抛异常会被 APScheduler
    吞掉并只留一行日志，不如自己接住并写进 schedule_runs 让用户看得见。"""
    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"}

    user_name = schedule.get("user_name") or "default"
    title = f"[定时] {schedule.get('title') or schedule_id[:8]}"

    session = await asyncio.to_thread(
        chat_store.create_session, user_name, schedule.get("mode") or "full", title
    )
    thread_id = session["id"]
    run_id = await asyncio.to_thread(schedule_store.start_run, schedule_id, thread_id)

    logger.info(
        "定时任务开始 id=%s run=%s trigger=%s thread=%s",
        schedule_id, run_id[:8], trigger, thread_id,
    )

    prompt = schedule.get("prompt") or ""
    await asyncio.to_thread(chat_store.save_message, thread_id, "user", prompt)

    try:
        from agents import agent_registry

        agent = agent_registry.get_agent()
        result = await agent.ainvoke(
            {"messages": [HumanMessage(content=prompt)]},
            config={
                "configurable": {"thread_id": thread_id},
                "recursion_limit": SCHEDULED_RECURSION_LIMIT,
            },
        )
    except Exception as exc:
        detail = f"{type(exc).__name__}: {exc}"
        logger.exception("定时任务失败 id=%s run=%s", schedule_id, run_id[:8])
        await asyncio.to_thread(
            chat_store.save_message, thread_id, "assistant", f"（任务执行失败）{detail}"
        )
        await asyncio.to_thread(schedule_store.finish_run, run_id, "failed", detail, None)
        return {"ok": False, "run_id": run_id, "session_id": thread_id, "error": detail}

    output = _extract_final_text(result)
    await asyncio.to_thread(chat_store.save_message, thread_id, "assistant", output)
    await asyncio.to_thread(schedule_store.finish_run, run_id, "success", None, output[:2000])

    # 计费：定时任务同样消耗额度，不记就是白跑
    input_tokens, output_tokens, source = _sum_usage(result.get("messages") or [])
    if input_tokens or output_tokens:
        meter = UsageMeter(
            user_name=user_name,
            session_id=thread_id,
            kind="scheduled",
            model=config.MODEL,
        )
        meter.add(input_tokens, output_tokens, source)
        meter.settle()

    logger.info(
        "定时任务完成 id=%s run=%s 输出长度=%d tokens=%d/%d",
        schedule_id, run_id[:8], len(output), input_tokens, output_tokens,
    )
    return {
        "ok": True,
        "run_id": run_id,
        "session_id": thread_id,
        "output": output,
        "input_tokens": input_tokens,
        "output_tokens": output_tokens,
    }


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"]