From 70b27159962e3cfce0048023ecdedcea5e97836a Mon Sep 17 00:00:00 2001
From: Xuwznln <18435084+Xuwznln@users.noreply.github.com>
Date: Tue, 14 Oct 2025 00:05:41 +0800
Subject: [PATCH] Fix workstation resource not tracking

---
 unilabos/devices/workstation/workstation_base.py |  1 -
 unilabos/ros/nodes/resource_tracker.py           | 10 ++++++++--
 unilabos/ros/utils/driver_creator.py             |  3 +++
 3 files changed, 11 insertions(+), 3 deletions(-)

diff --git a/unilabos/devices/workstation/workstation_base.py b/unilabos/devices/workstation/workstation_base.py
index 1988249f..97db1505 100644
--- a/unilabos/devices/workstation/workstation_base.py
+++ b/unilabos/devices/workstation/workstation_base.py
@@ -171,7 +171,6 @@ class WorkstationBase(ABC):
     def post_init(self, ros_node: ROS2WorkstationNode) -> None:
         # 初始化物料系统
         self._ros_node = ros_node
-        self._ros_node.update_resource([self.deck])
 
     def _build_resource_mappings(self, deck: Deck):
         """递归构建资源映射"""
diff --git a/unilabos/ros/nodes/resource_tracker.py b/unilabos/ros/nodes/resource_tracker.py
index 4ab31f0f..bd5b9184 100644
--- a/unilabos/ros/nodes/resource_tracker.py
+++ b/unilabos/ros/nodes/resource_tracker.py
@@ -1037,13 +1037,19 @@ class DeviceNodeResourceTracker(object):
     ) -> List[Tuple[Any, Any]]:
         res_list = []
         # print(resource, target_resource_cls_type, identifier_key, compare_value)
-        children = getattr(resource, "children", [])
+        children = []
+        if not isinstance(resource, dict):
+            children = getattr(resource, "children", [])
+        else:
+            children = resource.get("children")
+            if children is not None:
+                children = list(children.values()) if isinstance(children, dict) else children
         for child in children:
             res_list.extend(
                 self.loop_find_resource(child, target_resource_cls_type, identifier_key, compare_value, resource)
             )
         if issubclass(type(resource), target_resource_cls_type):
-            if target_resource_cls_type == dict:
+            if type(resource) == dict:
                 # 对于字典类型，直接检查 identifier_key
                 if identifier_key in resource:
                     if resource[identifier_key] == compare_value:
diff --git a/unilabos/ros/utils/driver_creator.py b/unilabos/ros/utils/driver_creator.py
index f72edf29..9481ce31 100644
--- a/unilabos/ros/utils/driver_creator.py
+++ b/unilabos/ros/utils/driver_creator.py
@@ -336,6 +336,9 @@ class WorkstationNodeCreator(DeviceClassCreator[T]):
         try:
             # 创建实例，额外补充一个给protocol node的字段，后面考虑取消
             data["children"] = self.children
+            for material_id, child in self.children.items():
+                if child["type"] != "device":
+                    self.resource_tracker.add_resource(self.children[material_id])
             deck_dict = data.get("deck")
             if deck_dict:
                 from pylabrobot.resources import Deck, Resource