vent_agent_2.0/runtime/business_tasks.py

"""持久化业务任务及结果投递；取数网关与结果接收平台各用独立连接。"""
import asyncio
from datetime import timedelta
import hashlib
import json
from pathlib import Path
import re
import time

import config
from agents.workspace import session_root
from db import business_store as store, chat_store, db_conn
from extensions.registry import load_extensions
from obs import get_logger
from runtime import business, execution
from runtime.business_models import RemoteFile, platform_result
from runtime.result_models import PlatformResultPayload

logger = get_logger(__name__)
_controller = None


def get_controller():
    if _controller is None:
        raise RuntimeError("业务任务执行器未启动，请在 RuntimeEnvironment(business_jobs=True) 内使用")
    return _controller


def submit(name, data, *, user_name="default", files=(), request_id=None, publish=False, source="manual"):
    spec = load_extensions().business(name)
    business.ensure_enabled(name)
    validated = spec.input_schema.model_validate(data).model_dump(mode="json")
    sources = [RemoteFile.model_validate(f).model_dump(mode="json") for f in files]
    needs_gateway = bool(sources or validated.get("use_gateway"))
    if needs_gateway and not spec.gateway.agent_code:
        raise ValueError("取数需要先配置本业务的网关调用身份")
    if needs_gateway and not config.get_bool("SHANXI_GATEWAY_ENABLED", False):
        raise ValueError("网关连接尚未启用")
    if publish and not spec.result_binding.agent_code:
        raise ValueError("结果推送需要配置本业务的结果平台身份")
    if publish and not config.get_bool("RESULT_PLATFORM_ENABLED", False):
        raise ValueError("结果接收平台尚未启用")
    return store.enqueue(name, user_name, validated, sources, request_id=request_id, source=source, publish=publish)[0]


def public_job(job):
    return {key: value for key, value in job.items() if key not in {"payload", "push_payload"}}


class BusinessTaskController:
    def __init__(self, gateway=None, *, publisher=None, executor=None):
        self.gateway = gateway
        self.publisher = publisher
        self.executor = executor or execution.execute
        self.catalog = load_extensions()
        self.tasks = []
        self.running_jobs = {}
        self.closed = False

    async def start(self):
        global _controller
        if _controller is not None:
            raise RuntimeError("业务任务执行器已经启动")
        await asyncio.to_thread(store.recover)
        for spec in self.catalog.business_agents.values():
            await asyncio.to_thread(store.ensure_state, spec.name)
        _controller = self
        for spec in self.catalog.business_agents.values():
            for index in range(spec.max_concurrency):
                self.tasks.append(asyncio.create_task(self._worker(spec.name), name=f"business-worker-{spec.name}-{index}"))
        if self.publisher:
            self.tasks.append(asyncio.create_task(self._publish_loop(), name="business-publisher"))

    async def close(self):
        global _controller
        self.closed = True
        for task in self.tasks:
            task.cancel()
        await asyncio.gather(*self.tasks, return_exceptions=True)
        self.tasks.clear()
        if _controller is self:
            _controller = None

    async def _worker(self, name):
        while not self.closed:
            try:
                await asyncio.to_thread(business.ensure_enabled, name)
                job = await asyncio.to_thread(store.claim, name)
                if job:
                    # 单次任务放在子 task，业务停止只取消该任务，worker 仍可恢复。
                    task = asyncio.create_task(self.run_job(job), name=f"business-job-{job['id']}")
                    self.running_jobs[job['id']] = task
                    try:
                        await task
                    except asyncio.CancelledError:
                        if self.closed:
                            raise
                    finally:
                        self.running_jobs.pop(job['id'], None)
            except business.BusinessStopped:
                pass
            except asyncio.CancelledError:
                raise
            except Exception as exc:
                logger.error("业务队列处理异常 agent=%s type=%s", name, type(exc).__name__)
            await asyncio.sleep(1)

    async def run_job(self, job):
        name, job_id = job["agent_name"], job["id"]
        with business.track(name):
            try:
                spec = self.catalog.business(name)
                await asyncio.to_thread(business.ensure_enabled, name)
                payload = json.loads(job["payload"])
                if not job.get("session_id"):
                    session = await asyncio.to_thread(chat_store.create_session, job["user_name"], "full", spec.display_name)
                    job["session_id"] = session["id"]
                    await asyncio.to_thread(store.update, job_id, session_id=session["id"])
                data = spec.input_schema.model_validate(payload["input"]).model_dump(mode="json")
                downloaded = []
                for raw in payload["files"]:
                    if self.gateway is None:
                        raise ValueError("网关未启用，不能下载平台文件")
                    source = RemoteFile.model_validate(raw)
                    filename = re.sub(r"[^\w.\-\u4e00-\u9fff]", "_", Path(source.filename).name)
                    key = hashlib.sha256(f"{source.file_id}\0{source.version}".encode()).hexdigest()[:12]
                    root = session_root(job["user_name"], job["session_id"])
                    target = root / "workspace" / "uploads" / f"{key}_{filename}"
                    metadata = await self.gateway.download(spec.gateway.agent_code, source, target)
                    await asyncio.to_thread(store.record_file, job_id, metadata)
                    downloaded.append({"file_id": source.file_id, "version": source.version,
                                       "path": f"/workspace/uploads/{target.name}"})
                # 文件路径以额外上下文传递，避免修改各业务输入 schema。
                message = "请执行本业务任务，按业务结果 schema 返回。\n输入：" + json.dumps(data, ensure_ascii=False)
                if downloaded:
                    if isinstance(data.get("materials"), list):
                        data["materials"] = [*data["materials"], *downloaded]
                        message = "请执行本业务任务，按业务结果 schema 返回。\n输入：" + json.dumps(data, ensure_ascii=False)
                    message += "\n已下载附件：" + json.dumps(downloaded, ensure_ascii=False)
                task_config = {"business_task_input": data, "business_task_id": job_id}
                if spec.skill_path:
                    # 大量时序已经通过服务端配置传给工具，不再作为聊天上下文重复展开。
                    summary = {key: data.get(key) for key in ('mine_code','mine_name','fan_id','fan_name','fan_pair','materials')}
                    summary.update(fan_state_count=len(data.get('fan_state_records',[])), flow_series_count=len(data.get('operating_series',[])))
                    message = ('请先读取 '+spec.skill_path+'，按 Skill 自主规划本次研判。'
                        '完整输入已经服务端校验并注入 business_task_input；先准备任务，按需使用网关取数。'
                        '文档解析、计算及报告生成都由你自行编写代码并在沙箱执行，根据真实输出调整计划；'
                        '最后提交脚本生成的 Markdown、研判 JSON 和成功执行ID。不要调用旧业务算法、使用历史报告或虚构原件事实。'
                        '\n本次输入摘要：'+json.dumps(summary,ensure_ascii=False))
                if job['source'] == 'test_pipeline':
                    result = await self._run_test_pipeline(spec, job, data)
                else:
                    result = await self.executor(agent_name=name, user_name=job["user_name"],
                        session_id=job["session_id"], message=message, kind="business",
                        graph_config={"recursion_limit": config.get_int(
                            'AGENT_EXPLORATION_RECURSION_LIMIT' if spec.skill_path else 'BUSINESS_RECURSION_LIMIT',
                            480 if spec.skill_path else 40),
                            "configurable": task_config})
                await self.complete(job, result)
            except asyncio.CancelledError:
                await asyncio.shield(asyncio.to_thread(store.update, job_id, status="cancelled", error="cancelled", finished_at=db_conn.fmt_dt()))
                raise
            except business.BusinessBusy:
                await asyncio.to_thread(store.update, job_id, status="queued", error=None)
            except business.BusinessStopped:
                await asyncio.to_thread(store.update, job_id, status="queued", error=None)
            except Exception as exc:
                error = type(exc).__name__
                if job['source'].startswith('test_') and isinstance(exc, ValueError):
                    error += ': ' + str(exc)[:300]
                await asyncio.to_thread(store.update, job_id, status="failed", error=error, finished_at=db_conn.fmt_dt())
                logger.error("业务任务失败 id=%s agent=%s type=%s", job_id, name, type(exc).__name__)

    async def _run_test_pipeline(self, spec, job, data):
        """测试业务工具本身；沿用生产准入、验证、归档与持久化队列。"""
        from runtime.engine import ExecutionResult
        if spec.test_runner is None:
            raise ValueError('该智能体未声明业务流程测试执行器，请选择完整智能体模式')
        # 纯规则样例无需模型计费配置；文档识别在真正调用模型前自行准入及计费。
        admission = await execution.admit(job['user_name'], job['session_id'], check_billing=False)
        lease = None
        try:
            lease = business.BusinessLease(spec.name)
            lease.bind()
            await asyncio.to_thread(store.bind_session, job['session_id'], spec.name)
            result = await spec.test_runner(data, job['user_name'], job['session_id'],
                gateway=self.gateway, job_id=job['id'],
                document_audit_dir=session_root(job['user_name'], job['session_id'])/'workspace/test_audits'/job['id'])
            text = f"# {result.get('result_title', spec.display_name)}\n\n{result.get('result','')}"
            await asyncio.to_thread(chat_store.save_message, job['session_id'], 'user', '智能体业务流程测试')
            await asyncio.to_thread(chat_store.save_message, job['session_id'], 'assistant', text, citations=result.get('sources', []))
            return ExecutionResult('', job['session_id'], 'success', text, '', {}, result.get('sources', []), [], result)
        finally:
            if lease:
                lease.release()
            admission.release()

    async def cancel_job(self, job_id):
        task = self.running_jobs.get(job_id)
        if task and not task.done():
            task.cancel()
            await asyncio.gather(task, return_exceptions=True)
        else:
            await asyncio.to_thread(db_conn.execute,
                "UPDATE business_jobs SET status='cancelled',finished_at=?,updated_at=? WHERE id=? AND status='queued'",
                (db_conn.fmt_dt(), db_conn.fmt_dt(), job_id))

    async def complete(self, job, result):
        fields = {"run_id": result.run_id, "error": result.error or None}
        if result.status != "success":
            fields.update(status=result.status, finished_at=None if result.status == "waiting_input" else db_conn.fmt_dt())
        else:
            spec = self.catalog.business(job["agent_name"])
            raw = result.structured_output
            if raw is None:
                if spec.skill_path:
                    raise ValueError('自主探索未完成正式提交，请查看工具反馈；需要通过报告提交工具归档后才能完成任务')
                try:
                    raw = json.loads(result.output)
                except json.JSONDecodeError as exc:
                    raise ValueError(
                        "业务智能体未交回有效的结构化结果，返回正文无法解析为 JSON；"
                        "请检查模型的思考模式及结果必填字段"
                    ) from exc
            if spec.result_validator:
                raw = await asyncio.to_thread(spec.result_validator, raw, job["user_name"], job["session_id"])
            checked = spec.result_schema.model_validate(raw)
            normalized = platform_result(checked.model_dump(mode="json"))
            source_input = spec.input_schema.model_validate(json.loads(job["payload"])["input"]).model_dump(mode="json")
            for key in ("mine_code", "mine_name"):
                if key in source_input and getattr(normalized, key) != source_input[key]:
                    raise ValueError("结果矿井身份与任务输入不一致")
            if normalized.platform_payload:
                exported = normalized.platform_payload
                if exported.taskId != job["id"] or (exported.mineCode, exported.mineName) != (normalized.mine_code, normalized.mine_name):
                    raise ValueError("结果平台报文身份与当前任务不一致")
            fields.update(status=normalized.status, result_json=checked.model_dump_json(),
                finished_at=normalized.platform_payload.checkedAt if normalized.platform_payload else db_conn.fmt_dt())
        await asyncio.to_thread(store.update, job["id"], **fields)
        if fields.get("status") == "completed" and json.loads(job["payload"])["publish"]:
            try:
                await self.queue_publish(await asyncio.to_thread(store.get, job["id"]))
            except Exception as exc:
                # 推送配置或映射错误单独记录，已经完成的业务结果仍保留。
                await asyncio.to_thread(store.update, job["id"], push_status="failed", error=type(exc).__name__)

    async def on_resume(self, session_id, result):
        rows = await asyncio.to_thread(db_conn.query,
            "SELECT * FROM business_jobs WHERE session_id = ? AND status = 'waiting_input'", (session_id,))
        for job in rows:
            try:
                await self.complete(job, result)
            except Exception as exc:
                await asyncio.to_thread(store.update, job["id"], status="failed", error=type(exc).__name__)

    async def _prepare_push(self, job):
        spec = self.catalog.business(job["agent_name"])
        if self.publisher is None:
            raise ValueError("结果接收平台尚未启用")
        result = json.loads(job["result_json"])
        if result["status"] != "completed":
            raise ValueError("未完成的业务结果不能作为正式研判上报")
        if not result.get("platform_payload"):
            raise ValueError("历史结果没有新格式报文，请重新执行研判")
        payload = PlatformResultPayload.model_validate(result["platform_payload"]).model_dump(mode="json")
        if payload["taskId"] != job["id"] or (payload["mineCode"], payload["mineName"]) != (result["mine_code"], result["mine_name"]):
            raise ValueError("推送结果身份与任务不一致")
        if payload["agentCode"] != spec.result_binding.agent_code or payload["agentName"] != spec.result_binding.agent_name:
            raise ValueError("结果平台身份已改变，不能改写历史结果")
        return await asyncio.to_thread(store.freeze_push, job["id"], payload, self.publisher.settings.url)

    async def queue_publish(self, job):
        if job["push_status"] == "pushed":
            return
        await self._prepare_push(job)
        await asyncio.to_thread(store.update, job["id"], push_status="pending", push_attempts=0, next_push_at=None, error=None)

    async def publish_job(self, job):
        with business.track(job["agent_name"]):
            await self._publish_job_once(job)

    async def _publish_job_once(self, job):
        name = job["agent_name"]
        started = time.monotonic()
        try:
            await asyncio.to_thread(business.ensure_enabled, name)
            if self.publisher is None:
                raise ValueError("结果接收平台尚未启用")
            payload = PlatformResultPayload.model_validate(json.loads(job["push_payload"])).model_dump(mode="json") if job.get("push_payload") else await self._prepare_push(job)
            # 推送重试沿用首次生成的身份和 taskId，禁止身份变化后误归属。
            spec = self.catalog.business(name)
            if (payload["agentCode"], payload["agentName"]) != (spec.result_binding.agent_code, spec.result_binding.agent_name):
                raise ValueError("平台身份已改变，请人工确认待推送记录")
            if await asyncio.to_thread(store.push_target, job["id"]) != self.publisher.settings.url:
                raise ValueError("结果接收平台地址已改变，不能改投历史结果")
            receipt = await self.publisher.push(payload)
            await asyncio.to_thread(store.update, job["id"], push_status="pushed", receipt_json=json.dumps(receipt, ensure_ascii=False), next_push_at=None, error=None)
            logger.info('研判结果推送成功', extra={'kv': {'event': 'result.push.completed',
                'task_id': job['id'], 'agent': name, 'duration_ms': round((time.monotonic() - started) * 1000)}})
        except business.BusinessStopped:
            return
        except asyncio.CancelledError:
            raise
        except Exception as exc:
            from extensions.result_platform.client import ResultPlatformError
            attempts = job["push_attempts"] + 1
            retryable = isinstance(exc, ResultPlatformError) and exc.retryable
            delay = min(2 ** min(attempts, 10), 600)
            if isinstance(exc, ResultPlatformError) and exc.retry_after is not None:
                delay = max(delay, exc.retry_after)
            next_at = db_conn.fmt_dt(db_conn.now() + timedelta(seconds=delay))
            await asyncio.to_thread(store.update, job["id"], push_status="pending" if retryable and attempts < 10 else "failed",
                                    push_attempts=attempts, next_push_at=next_at, error=type(exc).__name__)
            logger.warning('研判结果推送失败', exc_info=True, extra={'kv': {'event': 'result.push.failed',
                'task_id': job['id'], 'agent': name, 'error_type': type(exc).__name__, 'attempt': attempts,
                'retryable': retryable, 'next_push_at': next_at if retryable and attempts < 10 else None,
                'duration_ms': round((time.monotonic() - started) * 1000)}})

    async def _publish_loop(self):
        while not self.closed:
            try:
                for job in await asyncio.to_thread(store.pending_pushes, db_conn.fmt_dt()):
                    task = asyncio.create_task(self.publish_job(job), name=f"business-push-{job['id']}")
                    try:
                        await task
                    except asyncio.CancelledError:
                        if self.closed:
                            raise
            except asyncio.CancelledError:
                raise
            except Exception as exc:
                logger.error("推送队列异常 type=%s", type(exc).__name__)
            await asyncio.sleep(1)