front/xhs-agent-desktop/app/oms_agent_runner.py

"""OMS Agent 主循环：心跳、拉取任务、执行浏览器自动化。"""

from __future__ import annotations

import logging
import shutil
import tempfile
import time
import uuid
from pathlib import Path
from typing import Any, Callable
from urllib.parse import urlparse

import requests

from app.oms_client import OmsAgentClient, OmsClientError
from xhs_bridge.config import parse_human_behavior_config, parse_human_guard_config
from xhs_bridge.exceptions import HumanGuardError, XhsBridgeError
from xhs_bridge.service import XhsBridgeService

logger = logging.getLogger(__name__)

HEARTBEAT_INTERVAL_SEC = 30
POLL_INTERVAL_SEC = 5


def run_agent_loop(
    client: OmsAgentClient,
    service: XhsBridgeService,
    *,
    device_name: str | None = None,
    bridge_port: int = 9333,
    stop_flag: Callable[[], bool] | None = None,
) -> None:
    client.register(device_name=device_name, bridge_port=bridge_port)
    logger.info("Agent 已注册 deviceId=%s agentId=%s", client.device_id, client.agent_id)

    last_heartbeat = 0.0
    temp_dir = Path(tempfile.mkdtemp(prefix="xhs_agent_media_"))

    try:
        while not (stop_flag and stop_flag()):
            now = time.monotonic()
            if now - last_heartbeat >= HEARTBEAT_INTERVAL_SEC:
                _heartbeat(client, service)
                _sync_human_config(client, service)
                last_heartbeat = now

            tasks = client.claim_tasks(limit=1)
            if not tasks:
                time.sleep(POLL_INTERVAL_SEC)
                continue

            for task in tasks:
                _execute_task(client, service, task, client.base_url, temp_dir)
    finally:
        shutil.rmtree(temp_dir, ignore_errors=True)


def _heartbeat(client: OmsAgentClient, service: XhsBridgeService) -> None:
    status = service.status()
    client.heartbeat(
        engine_running=bool(status.get("bridge_running")),
        extension_connected=bool(status.get("extension_connected")),
        xhs_logged_in=bool(status.get("login_ok")),
        last_error=status.get("last_error") if isinstance(status.get("last_error"), str) else None,
    )


def _sync_human_config(client: OmsAgentClient, service: XhsBridgeService) -> None:
    try:
        pulled = client.pull_human_config(client.human_config_version)
    except OmsClientError as exc:
        logger.warning("拉取拟人配置失败: %s", exc)
        return

    if pulled.get("unchanged"):
        return

    behavior_raw = pulled.get("behavior")
    guard_raw = pulled.get("guard")
    if isinstance(behavior_raw, dict):
        service.apply_human_behavior_config(parse_human_behavior_config(behavior_raw))
    if isinstance(guard_raw, dict):
        service.apply_human_guard_config(parse_human_guard_config(guard_raw))

    version = pulled.get("version")
    if version is not None:
        client.set_human_config_version(int(version))
    logger.info("已同步 OMS 拟人配置 version=%s", client.human_config_version)


def _execute_task(
    client: OmsAgentClient,
    service: XhsBridgeService,
    task: dict[str, Any],
    oms_base_url: str,
    temp_dir: Path,
) -> None:
    task_id = int(task["id"])
    task_type = str(task.get("taskType") or "")
    tool_name = task.get("toolName")
    arguments = task.get("arguments") or {}
    if not isinstance(arguments, dict):
        arguments = {}

    client.report_task(task_id, status="running", progress_text="Agent 开始执行")
    try:
        result = _dispatch_task(service, task_type, tool_name, arguments, oms_base_url, temp_dir)
        client.report_task(
            task_id,
            status="succeeded",
            progress_text="执行完成",
            result=result,
        )
    except Exception as exc:
        logger.exception("任务 %s 失败", task_id)
        client.report_task(
            task_id,
            status="failed",
            progress_text=str(exc)[:500],
            error_message=str(exc)[:1000],
        )


def _dispatch_task(
    service: XhsBridgeService,
    task_type: str,
    tool_name: str | None,
    arguments: dict[str, Any],
    oms_base_url: str,
    temp_dir: Path,
) -> dict[str, Any]:
    if task_type == "engine_start":
        pid = service.start_engine()
        return {"pid": pid, "summary": f"引擎已启动 pid={pid}"}

    if task_type == "engine_stop":
        service.stop_engine()
        return {"summary": "引擎已停止"}

    if task_type == "browser_login":
        result = service.browser_login()
        if not result.get("logged_in"):
            raise XhsBridgeError(str(result.get("message") or "登录未完成"))
        return {"summary": str(result.get("message") or "登录成功"), **result}

    if task_type == "clear_cookies":
        result = service.clear_cookies(require_engine=True)
        return {"summary": "Cookie 已清除", **result}

    if task_type == "check_login":
        status = service.status()
        text = service.call_tool("check_login_status", {}, check_license=False)
        return {
            "summary": text,
            "login_ok": bool(status.get("login_ok")),
            "raw": text,
        }

    if not tool_name:
        raise XhsBridgeError(f"任务缺少 toolName: {task_type}")

    args = dict(arguments)
    if "images" in args and isinstance(args["images"], list):
        args["images"] = _resolve_media_paths(args["images"], oms_base_url, temp_dir)
    if "video" in args and isinstance(args["video"], str):
        args["video"] = _resolve_media_paths([args["video"]], oms_base_url, temp_dir)[0]

    def on_progress(msg: str) -> None:
        logger.info("[%s] %s", tool_name, msg)

    try:
        # OMS JWT 已鉴权，不再要求 LIKE AI 本地激活码
        text = service.call_tool(
            tool_name,
            args,
            check_license=False,
            on_progress=on_progress,
        )
    except HumanGuardError as exc:
        raise XhsBridgeError(f"拟人防护拦截: {exc}") from exc

    success = "❌" not in text and "失败" not in text
    if not success:
        raise XhsBridgeError(text)

    return {"summary": text[:500], "raw": text}


def _resolve_media_paths(urls: list[Any], oms_base_url: str, temp_dir: Path) -> list[str]:
    resolved: list[str] = []
    for item in urls:
        path = str(item or "").strip()
        if not path:
            continue
        if path.startswith(("http://", "https://")):
            resolved.append(_download_url(path, temp_dir))
        elif path.startswith("/"):
            resolved.append(_download_url(_absolute_url(oms_base_url, path), temp_dir))
        else:
            resolved.append(path)
    return resolved


def _absolute_url(oms_base_url: str, relative: str) -> str:
    origin = oms_base_url.rstrip("/")
    if origin.endswith("/admin"):
        origin = origin[: -len("/admin")]
    return origin + relative


def _download_url(url: str, temp_dir: Path) -> str:
    parsed = urlparse(url)
    suffix = Path(parsed.path).suffix or ".bin"
    target = temp_dir / f"{uuid.uuid4().hex}{suffix}"
    resp = requests.get(url, timeout=120)
    resp.raise_for_status()
    target.write_bytes(resp.content)
    return str(target.resolve())