run.py

import json
import logging

import psutil
from agents.agent import GlobalAgent
from flask import Flask, jsonify, request
from flask_cors import CORS
from flask_socketio import SocketIO

from task_manager import (
    ContextStore,
    TaskDiagnostics,
    TaskMonitor,
    TaskQueueManager,
    TaskScheduler,
    WSBridge,
)
from task_manager import register as register_hmi


logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
)
logger = logging.getLogger("master.run")

app = Flask(__name__, static_folder="assets")
CORS(app, resources={r"/*": {"origins": "*"}})
socketio = SocketIO(app, cors_allowed_origins="*")


master_agent = GlobalAgent(config_path="config.yaml")


# ──────────────────────────────────────────────────────────────────────
# Task Manager wiring（蒲）
# ──────────────────────────────────────────────────────────────────────
_tm_cfg = (master_agent.config.get("task_manager") or {})
_q_cfg = _tm_cfg.get("queue") or {}
_run_cfg = _tm_cfg.get("runtime") or {}
_hmi_cfg = _tm_cfg.get("hmi") or {}
_planner_cfg = _tm_cfg.get("planner") or {}

# 讓 scene profile 的 properties 能覆寫 runtime 預設
_scene_props = (master_agent.scene or {}).get("properties", {}) or {}
_run_cfg.setdefault("grasp_timeout_sec", _scene_props.get("grasp_timeout_sec", 8.0))
_run_cfg.setdefault("push_timeout_sec", _scene_props.get("push_timeout_sec", 6.0))
_run_cfg.setdefault("reorient_max_retry", _scene_props.get("reorient_max_retry", 3))

context_store = ContextStore(collaborator=master_agent.collaborator)
queue_mgr = TaskQueueManager(max_active=int(_q_cfg.get("max_active", 3)))
# diagnostics 需要先有 callback ref；先 None 之後 patch
diagnostics = TaskDiagnostics(
    collaborator=master_agent.collaborator,
    context_store=context_store,
    scheduler_action_callback=None,
)
monitor = TaskMonitor(
    collaborator=master_agent.collaborator,
    context_store=context_store,
    diagnostics=diagnostics,
)
scheduler = TaskScheduler(
    global_agent=master_agent,
    context_store=context_store,
    queue=queue_mgr,
    monitor=monitor,
    diagnostics=diagnostics,
    max_replan_attempts=int(_planner_cfg.get("max_replan_attempts", 3)),
)
# 把 scheduler 的 diagnostics callback 補上
diagnostics._scheduler_cb = scheduler.diagnostics_callback  # noqa: SLF001

# WebSocket bridge：訂閱 task_progress / task_faults → socket.io
ws_bridge = WSBridge(
    socketio=socketio,
    collaborator=master_agent.collaborator,
    progress_channel=(_tm_cfg.get("monitor") or {}).get("progress_channel", "task_progress"),
    fault_channel=(_tm_cfg.get("monitor") or {}).get("fault_channel", "task_faults"),
)
# monitor.on_progress 直接送一份給 socket.io（不必經 Redis 一圈）
monitor.on_progress = ws_bridge.push_progress  # type: ignore[assignment]

# 啟動 listener thread
monitor.start_listeners()
ws_bridge.start()

# 註冊 HTTP blueprint
register_hmi(
    app=app,
    scheduler=scheduler,
    monitor=monitor,
    diagnostics=diagnostics,
    context_store=context_store,
    url_prefix=str(_hmi_cfg.get("url_prefix", "") or ""),
)
logger.info("Task Manager wired up: queue=%d, mode=%s", queue_mgr.max_active, scheduler.mode.value)


def send_text_to_forntend(text):
    socketio.emit("text_update", {"data": text}, namespace="/")


@app.route("/system_status", methods=["GET"])
def system_status():
    """
    Get the system status.

    Returns:
        JSON response with system status
    """
    cpu_load = psutil.cpu_percent(interval=None)

    memory = psutil.virtual_memory()
    memory_usage = memory.percent

    return jsonify(
        {
            "cpu_load": round(cpu_load, 1),
            "memory_usage": round(memory_usage, 1),
        }
    )


@app.route("/robot_status", methods=["GET"])
def robot_status():
    """
    Get the status of all robots.

    Returns:
        JSON response with robot status
    """
    try:
        registered_robots = master_agent.collaborator.read_all_agents_info()
        registered_robots_status = []
        for robot_name, robot_info in registered_robots.items():
            registered_robots_status.append(
                {
                    "robot_name": robot_name,
                    "robot_state": json.loads(robot_info).get("robot_state"),
                }
            )
        return jsonify(registered_robots_status), 200
    except Exception as e:
        return jsonify({"error": "Internal server error", "details": str(e)}), 500


@app.route("/publish_task", methods=["POST", "GET"])
def publish_task():
    """
    Publish a task to the Redis channel.

    Request JSON format:
    {
        "task": "task_content"  # The task to be published
        "refresh": "true" # Boolean value, default is true, indicating whether to refresh the cached robot memory
    }

    Returns:
        JSON response with status or error message
    """
    if request.method == "GET":
        return jsonify({"statis": "success"}), 200
    try:
        data = request.get_json()
        if not data or "task" not in data:
            return jsonify({"status": "error", "message": "Invalid request - 'task' field required"}), 400
        if not isinstance(data["task"], list):
            data["task"] = [data["task"]]
        if "refresh" not in data:
            data["refresh"] = False

        task_id = data.get("task_id")
        for task in data["task"]:
            if not isinstance(task, str):
                return jsonify({"status": "error", "message": "Invalid task format - must be a string"}), 400
            subtask_list = master_agent.publish_global_task(
                task, data["refresh"], task_id
            )

        return (
            jsonify(
                {
                    "status": "success",
                    "message": "Task published successfully",
                    "data": subtask_list,
                }
            ),
            200,
        )

    except Exception as e:
        return jsonify({"status": "error", "message": f"Internal server error: {str(e)}"}), 500

if __name__ == "__main__":
    # Run the Flask app
    app.run(host="0.0.0.0", port=5000)