更新物料接口

unilabos/app/backend.py

@@ -1,14 +1,15 @@
 import threading
 
+from unilabos.ros.nodes.resource_tracker import ResourceTreeSet
 from unilabos.utils import logger
 
 
 # 根据选择的 backend 启动相应的功能
 def start_backend(
     backend: str,
-    devices_config: dict = {},
-    resources_config: list = [],
-    resources_edge_config: list = [],
+    devices_config: ResourceTreeSet,
+    resources_config: ResourceTreeSet,
+    resources_edge_config: list[dict] = [],
     graph=None,
     controllers_config: dict = {},
     bridges=[],

unilabos/app/main.py

@@ -7,9 +7,12 @@ import sys
 import threading
 import time
 from copy import deepcopy
+from typing import Dict, Any, List
 
+import networkx as nx
 import yaml
 
+from unilabos.ros.nodes.resource_tracker import ResourceTreeSet, ResourceDict
 
 # 首先添加项目根目录到路径
 current_dir = os.path.dirname(os.path.abspath(__file__))
@@ -225,6 +228,15 @@ def main():
     else:
         HTTPConfig.remote_addr = args_dict.get("addr", "")
 
+    # 设置BasicConfig参数
+    if args_dict.get("ak", ""):
+        BasicConfig.ak = args_dict.get("ak", "")
+        print_status("传入了ak参数，优先采用传入参数！", "info")
+    if args_dict.get("sk", ""):
+        BasicConfig.sk = args_dict.get("sk", "")
+        print_status("传入了sk参数，优先采用传入参数！", "info")
+
+    # 使用远程资源启动
     if args_dict["use_remote_resource"]:
         print_status("使用远程资源启动", "info")
         from unilabos.app.web import http_client
@@ -236,13 +248,6 @@ def main():
         else:
             print_status("远程资源不存在，本地将进行首次上报！", "info")
 
-    # 设置BasicConfig参数
-    if args_dict.get("ak", ""):
-        BasicConfig.ak = args_dict.get("ak", "")
-        print_status("传入了ak参数，优先采用传入参数！", "info")
-    if args_dict.get("sk", ""):
-        BasicConfig.sk = args_dict.get("sk", "")
-        print_status("传入了sk参数，优先采用传入参数！", "info")
     BasicConfig.working_dir = working_dir
     BasicConfig.is_host_mode = not args_dict.get("is_slave", False)
     BasicConfig.slave_no_host = args_dict.get("slave_no_host", False)
@@ -278,6 +283,10 @@ def main():
     if not BasicConfig.ak or not BasicConfig.sk:
         print_status("后续运行必须拥有一个实验室，请前往 https://uni-lab.bohrium.com 注册实验室！", "warning")
         os._exit(1)
+    graph: nx.Graph
+    resource_tree_set: ResourceTreeSet
+    resource_links: List[Dict[str, Any]]
+
     if args_dict["graph"] is None:
         request_startup_json = http_client.request_startup_json()
         if not request_startup_json:
@@ -287,58 +296,54 @@ def main():
             os._exit(1)
         else:
             print_status("联网获取设备加载文件成功", "info")
-        graph, data = read_node_link_json(request_startup_json)
+        graph, resource_tree_set, resource_links = read_node_link_json(request_startup_json)
     else:
         file_path = args_dict["graph"]
         if file_path.endswith(".json"):
-            graph, data = read_node_link_json(file_path)
+            graph, resource_tree_set, resource_links = read_node_link_json(file_path)
         else:
-            graph, data = read_graphml(file_path)
+            graph, resource_tree_set, resource_links = read_graphml(file_path)
     import unilabos.resources.graphio as graph_res
 
     graph_res.physical_setup_graph = graph
-    resource_edge_info = modify_to_backend_format(data["links"])
+    resource_edge_info = modify_to_backend_format(resource_links)
     materials = lab_registry.obtain_registry_resource_info()
     materials.extend(lab_registry.obtain_registry_device_info())
     materials = {k["id"]: k for k in materials}
-    nodes = {k["id"]: k for k in data["nodes"]}
+    # 从 ResourceTreeSet 中获取节点信息
+    nodes = {node.res_content.id: node.res_content for node in resource_tree_set.all_nodes}
     edge_info = len(resource_edge_info)
     for ind, i in enumerate(resource_edge_info[::-1]):
-        source_node = nodes[i["source"]]
-        target_node = nodes[i["target"]]
+        source_node: ResourceDict = nodes[i["source"]]
+        target_node: ResourceDict = nodes[i["target"]]
         source_handle = i["sourceHandle"]
         target_handle = i["targetHandle"]
         source_handler_keys = [
-            h["handler_key"] for h in materials[source_node["class"]]["handles"] if h["io_type"] == "source"
+            h["handler_key"] for h in materials[source_node.klass]["handles"] if h["io_type"] == "source"
         ]
         target_handler_keys = [
-            h["handler_key"] for h in materials[target_node["class"]]["handles"] if h["io_type"] == "target"
+            h["handler_key"] for h in materials[target_node.klass]["handles"] if h["io_type"] == "target"
         ]
         if source_handle not in source_handler_keys:
             print_status(
-                f"节点 {source_node['id']} 的source端点 {source_handle} 不存在，请检查，支持的端点 {source_handler_keys}",
+                f"节点 {source_node.id} 的source端点 {source_handle} 不存在，请检查，支持的端点 {source_handler_keys}",
                 "error",
             )
             resource_edge_info.pop(edge_info - ind - 1)
             continue
         if target_handle not in target_handler_keys:
             print_status(
-                f"节点 {target_node['id']} 的target端点 {target_handle} 不存在，请检查，支持的端点 {target_handler_keys}",
+                f"节点 {target_node.id} 的target端点 {target_handle} 不存在，请检查，支持的端点 {target_handler_keys}",
                 "error",
             )
             resource_edge_info.pop(edge_info - ind - 1)
             continue
 
-    devices_and_resources = dict_from_graph(graph_res.physical_setup_graph)
-    # args_dict["resources_config"] = initialize_resources(list(deepcopy(devices_and_resources).values()))
-    args_dict["resources_config"] = list(devices_and_resources.values())
-    args_dict["devices_config"] = dict_to_nested_dict(deepcopy(devices_and_resources), devices_only=False)
+    # 使用 ResourceTreeSet 代替 list
+    args_dict["resources_config"] = resource_tree_set
+    args_dict["devices_config"] = resource_tree_set
     args_dict["graph"] = graph_res.physical_setup_graph
 
-    print_status(f"{len(args_dict['resources_config'])} Resources loaded:", "info")
-    for i in args_dict["resources_config"]:
-        print_status(f"DeviceId: {i['id']}, Class: {i['class']}", "info")
-
     if BasicConfig.upload_registry:
         # 设备注册到服务端 - 需要 ak 和 sk
         if args_dict.get("ak") and args_dict.get("sk"):
@@ -351,9 +356,7 @@ def main():
         else:
             print_status("未提供 ak 和 sk，跳过设备注册", "info")
     else:
-        print_status(
-            "本次启动注册表不报送云端，如果您需要联网调试，请在启动命令增加--upload_registry", "warning"
-        )
+        print_status("本次启动注册表不报送云端，如果您需要联网调试，请在启动命令增加--upload_registry", "warning")
 
     if args_dict["controllers"] is not None:
         args_dict["controllers_config"] = yaml.safe_load(open(args_dict["controllers"], encoding="utf-8"))
@@ -383,13 +386,16 @@ def main():
     # web visiualize 2D
     if args_dict["visual"] != "disable":
         enable_rviz = args_dict["visual"] == "rviz"
+        devices_and_resources = dict_from_graph(graph_res.physical_setup_graph)
         if devices_and_resources is not None:
             from unilabos.device_mesh.resource_visalization import (
                 ResourceVisualization,
             )  # 此处开启后，logger会变更为INFO，有需要请调整
 
             resource_visualization = ResourceVisualization(
-                devices_and_resources, args_dict["resources_config"], enable_rviz=enable_rviz
+                devices_and_resources,
+                [n.res_content for n in args_dict["resources_config"].all_nodes],  # type: ignore  # FIXME
+                enable_rviz=enable_rviz,
             )
             args_dict["resources_mesh_config"] = resource_visualization.resource_model
             start_backend(**args_dict)

unilabos/app/register.py

@@ -1,11 +1,6 @@
-import argparse
 import json
 import time
 
-from unilabos.config.config import BasicConfig
-from unilabos.registry.registry import build_registry
-
-from unilabos.app.main import load_config_from_file
 from unilabos.utils.log import logger
 from unilabos.utils.type_check import TypeEncoder
 

unilabos/app/web/client.py

@@ -9,6 +9,7 @@ import os
 from typing import List, Dict, Any, Optional
 
 import requests
+from unilabos.ros.nodes.resource_tracker import ResourceTreeSet
 from unilabos.utils.log import info
 from unilabos.config.config import HTTPConfig, BasicConfig
 from unilabos.utils import logger
@@ -46,7 +47,7 @@ class HTTPClient:
             Response: API响应对象
         """
         response = requests.post(
-            f"{self.remote_addr}/lab/material/edge",
+            f"{self.remote_addr}/edge/material/edge",
             json={
                 "edges": resources,
             },
@@ -61,6 +62,81 @@ class HTTPClient:
             logger.error(f"添加物料关系失败: {response.status_code}, {response.text}")
         return response
 
+    def resource_tree_add(self, resources: ResourceTreeSet, mount_uuid: str, first_add: bool) -> Dict[str, str]:
+        """
+        添加资源
+
+        Args:
+            resources: 要添加的资源树集合（ResourceTreeSet）
+            mount_uuid: 要挂载的资源的uuid
+            first_add: 是否为首次添加资源，可以是host也可以是slave来的
+        Returns:
+            Dict[str, str]: 旧UUID到新UUID的映射关系 {old_uuid: new_uuid}
+        """
+        # 从序列化数据中提取所有节点的UUID（保存旧UUID）
+        old_uuids = {n.res_content.uuid: n for n in resources.all_nodes}
+        if not self.initialized or first_add:
+            self.initialized = True
+            info(f"首次添加资源，当前远程地址: {self.remote_addr}")
+            response = requests.post(
+                f"{self.remote_addr}/edge/material",
+                json={"nodes": [x for xs in resources.dump() for x in xs], "mount_uuid": mount_uuid},
+                headers={"Authorization": f"Lab {self.auth}"},
+                timeout=100,
+            )
+        else:
+            response = requests.put(
+                f"{self.remote_addr}/edge/material",
+                json={"nodes": [x for xs in resources.dump() for x in xs], "mount_uuid": mount_uuid},
+                headers={"Authorization": f"Lab {self.auth}"},
+                timeout=100,
+            )
+
+        # 处理响应，构建UUID映射
+        uuid_mapping = {}
+        if response.status_code == 200:
+            res = response.json()
+            if "code" in res and res["code"] != 0:
+                logger.error(f"添加物料失败: {response.text}")
+            else:
+                data = res["data"]
+                for i in data:
+                    uuid_mapping[i["uuid"]] = i["cloud_uuid"]
+        else:
+            logger.error(f"添加物料失败: {response.text}")
+        for u, n in old_uuids.items():
+            if u in uuid_mapping:
+                n.res_content.uuid = uuid_mapping[u]
+            else:
+                logger.warning(f"资源UUID未更新: {u}")
+        return uuid_mapping
+
+    def resource_tree_get(self, uuid_list: List[str], with_children: bool) -> List[Dict[str, Any]]:
+        """
+        添加资源
+
+        Args:
+            uuid_list: List[str]
+        Returns:
+            Dict[str, str]: 旧UUID到新UUID的映射关系 {old_uuid: new_uuid}
+        """
+        response = requests.post(
+            f"{self.remote_addr}/edge/material/query",
+            json={"uuids": uuid_list, "with_children": with_children},
+            headers={"Authorization": f"Lab {self.auth}"},
+            timeout=100,
+        )
+        if response.status_code == 200:
+            res = response.json()
+            if "code" in res and res["code"] != 0:
+                logger.error(f"查询物料失败: {response.text}")
+            else:
+                data = res["data"]["nodes"]
+                return data
+        else:
+            logger.error(f"查询物料失败: {response.text}")
+        return []
+
     def resource_add(self, resources: List[Dict[str, Any]]) -> requests.Response:
         """
         添加资源

unilabos/app/ws_client.py

@@ -19,9 +19,12 @@ import websockets
 import ssl as ssl_module
 from queue import Queue, Empty
 from dataclasses import dataclass, field
-from typing import Optional, Dict, Any, Callable, List, Set
+from typing import Optional, Dict, Any, List
 from urllib.parse import urlparse
 from enum import Enum
+
+from jedi.inference.gradual.typing import TypedDict
+
 from unilabos.app.model import JobAddReq
 from unilabos.ros.nodes.presets.host_node import HostNode
 from unilabos.utils.type_check import serialize_result_info
@@ -96,6 +99,14 @@ class WebSocketMessage:
     timestamp: float = field(default_factory=time.time)
 
 
+class WSResourceChatData(TypedDict):
+    uuid: str
+    device_uuid: str
+    device_id: str
+    device_old_uuid: str
+    device_old_id: str
+
+
 class DeviceActionManager:
     """设备动作管理器 - 管理每个device_action_key的任务队列"""
 
@@ -543,7 +554,7 @@ class MessageProcessor:
     async def _process_message(self, data: Dict[str, Any]):
         """处理收到的消息"""
         message_type = data.get("action", "")
-        message_data = data.get("data", {})
+        message_data = data.get("data")
 
         logger.debug(f"[MessageProcessor] Processing message: {message_type}")
 
@@ -556,8 +567,12 @@ class MessageProcessor:
                 await self._handle_job_start(message_data)
             elif message_type == "cancel_action" or message_type == "cancel_task":
                 await self._handle_cancel_action(message_data)
-            elif message_type == "":
-                return
+            elif message_type == "add_material":
+                await self._handle_resource_tree_update(message_data, "add")
+            elif message_type == "update_material":
+                await self._handle_resource_tree_update(message_data, "update")
+            elif message_type == "remove_material":
+                await self._handle_resource_tree_update(message_data, "remove")
             else:
                 logger.debug(f"[MessageProcessor] Unknown message type: {message_type}")
 
@@ -574,6 +589,7 @@ class MessageProcessor:
     async def _handle_query_action_state(self, data: Dict[str, Any]):
         """处理query_action_state消息"""
         device_id = data.get("device_id", "")
+        device_uuid = data.get("device_uuid", "")
         action_name = data.get("action_name", "")
         task_id = data.get("task_id", "")
         job_id = data.get("job_id", "")
@@ -760,6 +776,92 @@ class MessageProcessor:
         else:
             logger.warning("[MessageProcessor] Cancel request missing both task_id and job_id")
 
+    async def _handle_resource_tree_update(self, resource_uuid_list: List[WSResourceChatData], action: str):
+        """处理资源树更新消息（add_material/update_material/remove_material）"""
+        if not resource_uuid_list:
+            return
+
+        # 按device_id和action分组
+        # device_action_groups: {(device_id, action): [uuid_list]}
+        device_action_groups = {}
+
+        for item in resource_uuid_list:
+            device_id = item["device_id"]
+            if not device_id:
+                device_id = "host_node"
+
+            # 特殊处理update action: 检查是否设备迁移
+            if action == "update":
+                device_old_id = item.get("device_old_id", "")
+                if not device_old_id:
+                    device_old_id = "host_node"
+
+                # 设备迁移：device_id != device_old_id
+                if device_id != device_old_id:
+                    # 给旧设备发送remove
+                    key_remove = (device_old_id, "remove")
+                    if key_remove not in device_action_groups:
+                        device_action_groups[key_remove] = []
+                    device_action_groups[key_remove].append(item["uuid"])
+
+                    # 给新设备发送add
+                    key_add = (device_id, "add")
+                    if key_add not in device_action_groups:
+                        device_action_groups[key_add] = []
+                    device_action_groups[key_add].append(item["uuid"])
+
+                    logger.info(
+                        f"[MessageProcessor] Resource migrated: {item['uuid'][:8]} from {device_old_id} to {device_id}"
+                    )
+                else:
+                    # 正常update
+                    key = (device_id, "update")
+                    if key not in device_action_groups:
+                        device_action_groups[key] = []
+                    device_action_groups[key].append(item["uuid"])
+            else:
+                # add或remove action，直接分组
+                key = (device_id, action)
+                if key not in device_action_groups:
+                    device_action_groups[key] = []
+                device_action_groups[key].append(item["uuid"])
+
+        logger.info(f"触发物料更新 {action} 分组数量: {len(device_action_groups)}, 总数量: {len(resource_uuid_list)}")
+
+        # 为每个(device_id, action)创建独立的更新线程
+        for (device_id, actual_action), items in device_action_groups.items():
+            logger.info(f"设备 {device_id} 物料更新 {actual_action} 数量: {len(items)}")
+
+            def _notify_resource_tree(dev_id, act, item_list):
+                try:
+                    host_node = HostNode.get_instance(timeout=5)
+                    if not host_node:
+                        logger.error(f"[MessageProcessor] HostNode instance not available for {act}")
+                        return
+
+                    success = host_node.notify_resource_tree_update(dev_id, act, item_list)
+
+                    if success:
+                        logger.info(
+                            f"[MessageProcessor] Resource tree {act} completed for device {dev_id}, "
+                            f"items: {len(item_list)}"
+                        )
+                    else:
+                        logger.warning(f"[MessageProcessor] Resource tree {act} failed for device {dev_id}")
+
+                except Exception as e:
+                    logger.error(f"[MessageProcessor] Error in resource tree {act} for device {dev_id}: {str(e)}")
+                    logger.error(traceback.format_exc())
+
+            # 在新线程中执行通知
+            thread = threading.Thread(
+                target=_notify_resource_tree,
+                args=(device_id, actual_action, items),
+                daemon=True,
+                name=f"ResourceTreeUpdate-{actual_action}-{device_id}",
+            )
+            thread.start()
+
     async def _send_action_state_response(
         self, device_id: str, action_name: str, task_id: str, job_id: str, typ: str, free: bool, need_more: int
     ):
@@ -1008,6 +1110,8 @@ class WebSocketClient(BaseCommunicationClient):
 
         # 构建WebSocket URL
         self.websocket_url = self._build_websocket_url()
+        if not self.websocket_url:
+            self.websocket_url = ""  # 默认空字符串，避免None
 
         # 两个核心线程
         self.message_processor = MessageProcessor(self.websocket_url, self.send_queue, self.device_manager)