vent_agent_2.0/runtime/execution.py

"""聊天、审批恢复和定时执行共用的准入与运行装配。"""
from __future__ import annotations

import asyncio
from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from runtime.engine import Execution, ExecutionResult

from runtime import concurrency
from billing import guard
from db import chat_store


@dataclass
class Admission:
    session: concurrency.SessionLease
    slot: concurrency.TaskSlot
    business_lease: object | None = None

    def release(self):
        if self.business_lease is not None:
            self.business_lease.release()
        self.slot.release()
        self.session.release()


async def admit(user_name: str, session_id: str, *, check_billing: bool = True) -> Admission:
    session = concurrency.acquire_session(session_id)
    slot = None
    try:
        slot = concurrency.acquire_task_slot(user_name)
        if not await asyncio.to_thread(chat_store.get_session_for_user, session_id, user_name):
            raise concurrency.SessionMissingError("会话不存在")
        if check_billing:
            await asyncio.to_thread(guard.precheck, user_name)
        return Admission(session, slot)
    except BaseException:
        if slot:
            slot.release()
        session.release()
        raise


async def build_execution(admission: Admission, *, agent=None, agent_name: str = "main", **kwargs) -> Execution:
    from agents import agent_registry
    from runtime.engine import Execution, ExecutionRequest
    try:
        from db import business_store
        if agent_name != "main":
            from extensions.registry import load_extensions
            load_extensions().business(agent_name)
        await asyncio.to_thread(business_store.bind_session, kwargs["session_id"], agent_name)
        if agent_name != "main":
            from runtime.business import BusinessLease
            from extensions.registry import load_extensions
            business = load_extensions().business(agent_name)
            admission.business_lease = BusinessLease(agent_name)
            kwargs.update(agent_name=agent_name, agent_cn_name=business.display_name,
                          business_lease=admission.business_lease, result_validator=business.result_validator)
        if agent is None:
            agent = (await asyncio.to_thread(agent_registry.get_agent) if agent_name == "main"
                     else await asyncio.to_thread(agent_registry.get_agent, agent_name))
        graph_config = kwargs["config"]
        if agent_name != 'main' and business.skill_path:
            import config
            graph_config.setdefault('recursion_limit', config.get_int('AGENT_EXPLORATION_RECURSION_LIMIT', 480))
        graph_config.setdefault("configurable", {}).update(
            workspace_user=kwargs["user_name"], workspace_session=kwargs["session_id"],
            business_agent=agent_name,
        )
        return Execution(ExecutionRequest(agent=agent, slot=admission.slot,
                         session_lease=admission.session, **kwargs))
    except BaseException:
        admission.release()
        raise


async def execute(*, user_name: str, session_id: str, message: str, agent=None,
                  agent_name: str = "main", kind: str = "batch",
                  graph_config: dict | None = None, trace_id: str = "") -> ExecutionResult:
    """非 HTTP 入口：对已有且归属当前用户的会话执行一次任务，返回完整结果。

    独立批量任务应各自创建会话；额度与并发限制和聊天入口一致。
    """
    from obs import bind_trace_id, unbind_trace_id, get_trace_id, log_context
    token = bind_trace_id(trace_id or get_trace_id())
    try:
        with log_context(session_id=session_id, user_name=user_name, kind=kind):
            from copy import deepcopy
            admission = await admit(user_name, session_id)
            try:
                options = deepcopy(graph_config) if graph_config is not None else {}
                options.setdefault("configurable", {})["thread_id"] = session_id
                execution = await build_execution(admission, agent=agent, agent_name=agent_name, session_id=session_id,
                    thread_id=session_id, user_name=user_name, message=message, original_user_message=message,
                    kind=kind, config=options, trace_id=get_trace_id())
                from runtime.interrupts import pending_interrupts, ResumeError
                if await pending_interrupts(execution.req.agent, options):
                    raise ResumeError("pending_interrupt", "请先处理当前待确认请求")
                return await execution.run()
            finally:
                admission.release()
    finally:
        unbind_trace_id(token)