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

    def release(self):
        self.slot.release()
        self.session.release()


async def admit(user_name: str, session_id: str) -> 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("会话不存在")
        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, **kwargs) -> Execution:
    from agents import agent_registry
    from runtime.engine import Execution, ExecutionRequest
    try:
        if agent is None:
            agent = await asyncio.to_thread(agent_registry.get_agent)
        graph_config = kwargs["config"]
        graph_config.setdefault("configurable", {}).update(
            workspace_user=kwargs["user_name"], workspace_session=kwargs["session_id"],
        )
        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,
                  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, 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)