forgotten context managers

dimos/agents2/skills/osm.py

@@ -22,13 +22,14 @@
 from dimos.protocol.skill.skill import SkillContainer, skill
 from dimos.robot.robot import Robot
 from dimos.utils.logging_config import setup_logger
+from dimos.core.resource import Resource
 
 from reactivex.disposable import CompositeDisposable
 
 logger = setup_logger(__file__)
 
 
-class OsmSkillContainer(SkillContainer):
+class OsmSkillContainer(SkillContainer, Resource):
     _robot: Robot
     _disposables: CompositeDisposable
     _latest_location: Optional[LatLon]
@@ -45,15 +46,13 @@ def __init__(self, robot: Robot, position_stream: Observable[LatLon]):
         self._current_location_map = CurrentLocationMap(QwenVlModel())
         self._started = False
 
-    def __enter__(self) -> "OsmSkillContainer":
+    def start(self):
         self._started = True
         self._disposables.add(self._position_stream.subscribe(self._on_gps_location))
-        return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
+    def stop(self):
         self._disposables.dispose()
-        self.stop()
-        return False
+        super().stop()
 
     def _on_gps_location(self, location: LatLon) -> None:
         self._latest_location = location

dimos/agents2/skills/ros_navigation.py

@@ -19,6 +19,7 @@
 from dimos.protocol.skill.skill import SkillContainer, skill
 from dimos.utils.logging_config import setup_logger
 from dimos.utils.transform_utils import euler_to_quaternion
+from dimos.core.resource import Resource
 
 from typing import TYPE_CHECKING
 
@@ -28,7 +29,7 @@
 logger = setup_logger(__file__)
 
 
-class RosNavigation(SkillContainer):
+class RosNavigation(SkillContainer, Resource):
     _robot: "UnitreeG1"
     _started: bool
 
@@ -37,13 +38,11 @@ def __init__(self, robot: "UnitreeG1"):
         self._similarity_threshold = 0.23
         self._started = False
 
-    def __enter__(self) -> "RosNavigation":
+    def start(self) -> None:
         self._started = True
-        return self
 
-    def __exit__(self, exc_type, exc_val, exc_tb):
-        self._close_module()
-        return False
+    def stop(self) -> None:
+        super().stop()
 
     @skill()
     def navigate_with_text(self, query: str) -> str:

dimos/mapping/osm/demo_osm.py

@@ -51,7 +51,7 @@ def start(self) -> None:
         self._robot = FakeRobot()
         self._agent = Agent(system_prompt=SYSTEM_PROMPT)
         self._osm_skill_container = OsmSkillContainer(self._robot, _get_fake_location())
-        self._osm_skill_container.__enter__()
+        self._osm_skill_container.start()
         self._agent.register_skills(self._osm_skill_container)
         self._agent.register_skills(HumanInput())
         self._agent.run_implicit_skill("human")
@@ -64,7 +64,7 @@ def stop(self) -> None:
         if self._robot_debugger:
             self._robot_debugger.stop()
         if self._osm_skill_container:
-            self._osm_skill_container.__exit__(None, None, None)
+            self._osm_skill_container.stop()
         if self._agent:
             self._agent.stop()
 

dimos/robot/unitree_webrtc/unitree_g1.py

@@ -201,6 +201,7 @@ def __init__(
         self.joystick = None
         self.ros_bridge = None
         self.camera = None
+        self._ros_nav = None
         self._setup_directories()
 
     def _setup_directories(self):
@@ -223,7 +224,7 @@ def _setup_directories(self):
         os.makedirs(self.db_path, exist_ok=True)
 
     def _deploy_detection(self, goto):
-        detection = self.dimos.deploy(
+        detection = self._dimos.deploy(
             ObjectDBModule, goto=goto, camera_info=zed.CameraInfo.SingleWebcam
         )
 
@@ -269,7 +270,7 @@ def start(self):
         if self.enable_ros_bridge:
             self._deploy_ros_bridge()
 
-        self.nav = self.dimos.deploy(NavigationModule)
+        self.nav = self._dimos.deploy(NavigationModule)
         self.nav.goal_reached.transport = core.LCMTransport("/goal_reached", Bool)
         self.nav.goal_pose.transport = core.LCMTransport("/goal_pose", PoseStamped)
         self.nav.cancel_goal.transport = core.LCMTransport("/cancel_goal", Bool)
@@ -294,16 +295,16 @@ def start(self):
         g1_skills = UnitreeG1SkillContainer(robot=self)
         agent.register_skills(g1_skills)
 
-        human_input = self.dimos.deploy(HumanInput)
+        human_input = self._dimos.deploy(HumanInput)
         agent.register_skills(human_input)
 
         if self.enable_perception:
             agent.register_skills(self.detection)
 
         # Register ROS navigation
-        ros_nav = RosNavigation(self)
-        ros_nav.__enter__()
-        agent.register_skills(ros_nav)
+        self._ros_nav = RosNavigation(self)
+        self._ros_nav.start()
+        agent.register_skills(self._ros_nav)
 
         agent.run_implicit_skill("human")
         agent.start()
@@ -320,6 +321,8 @@ def start(self):
 
     def stop(self) -> None:
         self._dimos.stop()
+        if self._ros_nav:
+            self._ros_nav.stop()
         self.lcm.stop()
 
     def _deploy_connection(self):
@@ -373,7 +376,7 @@ def _deploy_visualization(self):
         self.foxglove_bridge.start()
 
     def _deploy_perception(self):
-        self.spatial_memory_module = self.dimos.deploy(
+        self.spatial_memory_module = self._dimos.deploy(
             SpatialMemory,
             collection_name=self.spatial_memory_collection,
             db_path=self.db_path,

pyproject.toml

@@ -174,6 +174,7 @@ dev = [
     "pytest-asyncio==0.26.0",
     "pytest-mock==3.15.0",
     "pytest-env==1.1.5",
+    "pytest-timeout==2.4.0",
     "textual==3.7.1",
     "requests-mock==1.12.1",
 ]