memory2 StreamModules, cleanup

dimos/agents/agent_test_runner.py

@@ -49,8 +49,8 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(Disposable(self.agent.subscribe(self._on_agent_message)))
-        self._disposables.add(Disposable(self.agent_idle.subscribe(self._on_agent_idle)))
+        self.register_disposable(Disposable(self.agent.subscribe(self._on_agent_message)))
+        self.register_disposable(Disposable(self.agent_idle.subscribe(self._on_agent_idle)))
         # Signal that subscription is ready
         self._subscription_ready.set()
 

dimos/agents/mcp/mcp_client.py

@@ -168,7 +168,7 @@ def start(self) -> None:
         def _on_human_input(string: str) -> None:
             self._message_queue.put(HumanMessage(content=string))
 
-        self._disposables.add(Disposable(self.human_input.subscribe(_on_human_input)))
+        self.register_disposable(Disposable(self.human_input.subscribe(_on_human_input)))
 
     @rpc
     def on_system_modules(self, _modules: list[RPCClient]) -> None:

dimos/agents/skills/demo_robot.py

@@ -25,7 +25,7 @@ class DemoRobot(Module):
 
     def start(self) -> None:
         super().start()
-        self._disposables.add(interval(1.0).subscribe(lambda _: self._publish_gps_location()))
+        self.register_disposable(interval(1.0).subscribe(lambda _: self._publish_gps_location()))
 
     def stop(self) -> None:
         super().stop()

dimos/agents/skills/google_maps_skill_container.py

@@ -15,6 +15,8 @@
 import json
 from typing import Any
 
+from reactivex.disposable import Disposable
+
 from dimos.agents.annotation import skill
 from dimos.core.core import rpc
 from dimos.core.module import Module
@@ -49,7 +51,7 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(self.gps_location.subscribe(self._on_gps_location))  # type: ignore[arg-type]
+        self.register_disposable(Disposable(self.gps_location.subscribe(self._on_gps_location)))
 
     @rpc
     def stop(self) -> None:

dimos/agents/skills/gps_nav_skill.py

@@ -14,6 +14,8 @@
 
 import json
 
+from reactivex.disposable import Disposable
+
 from dimos.agents.annotation import skill
 from dimos.core.core import rpc
 from dimos.core.module import Module
@@ -37,7 +39,7 @@ class GpsNavSkillContainer(Module):
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(self.gps_location.subscribe(self._on_gps_location))  # type: ignore[arg-type]
+        self.register_disposable(Disposable(self.gps_location.subscribe(self._on_gps_location)))
 
     @rpc
     def stop(self) -> None:

dimos/agents/skills/navigation.py

@@ -68,8 +68,8 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(Disposable(self.color_image.subscribe(self._on_color_image)))
-        self._disposables.add(Disposable(self.odom.subscribe(self._on_odom)))
+        self.register_disposable(Disposable(self.color_image.subscribe(self._on_color_image)))
+        self.register_disposable(Disposable(self.odom.subscribe(self._on_odom)))
         self._skill_started = True
 
     @rpc

dimos/agents/skills/osm.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 
+from reactivex.disposable import Disposable
+
 from dimos.agents.annotation import skill
 from dimos.core.module import Module
 from dimos.core.stream import In
@@ -39,7 +41,7 @@ def __init__(self) -> None:
     def start(self) -> None:
         super().start()
         if hasattr(self.gps_location, "subscribe"):
-            self._disposables.add(self.gps_location.subscribe(self._on_gps_location))  # type: ignore[arg-type]
+            self.register_disposable(Disposable(self.gps_location.subscribe(self._on_gps_location)))
         else:
             logger.warning(
                 "OsmSkill: gps_location stream does not support direct subscribe (RemoteIn)"

dimos/agents/skills/person_follow.py

@@ -93,9 +93,9 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(Disposable(self.color_image.subscribe(self._on_color_image)))
+        self.register_disposable(Disposable(self.color_image.subscribe(self._on_color_image)))
         if self.config.use_3d_navigation:
-            self._disposables.add(Disposable(self.global_map.subscribe(self._on_pointcloud)))
+            self.register_disposable(Disposable(self.global_map.subscribe(self._on_pointcloud)))
 
     @rpc
     def stop(self) -> None:

dimos/agents/vlm_agent.py

@@ -16,6 +16,7 @@
 
 from langchain.chat_models import init_chat_model
 from langchain_core.messages import AIMessage, HumanMessage, SystemMessage
+from reactivex.disposable import Disposable
 
 from dimos.agents.system_prompt import SYSTEM_PROMPT
 from dimos.core.core import rpc
@@ -60,8 +61,8 @@ def __init__(self, **kwargs: Any) -> None:
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(self.color_image.subscribe(self._on_image))  # type: ignore[arg-type]
-        self._disposables.add(self.query_stream.subscribe(self._on_query))  # type: ignore[arg-type]
+        self.register_disposable(Disposable(self.color_image.subscribe(self._on_image)))
+        self.register_disposable(Disposable(self.query_stream.subscribe(self._on_query)))
 
     @rpc
     def stop(self) -> None:

dimos/agents/vlm_stream_tester.py

@@ -16,6 +16,7 @@
 import time
 
 from langchain_core.messages import AIMessage, HumanMessage
+from reactivex.disposable import Disposable
 
 from dimos.core.core import rpc
 from dimos.core.module import Module
@@ -62,8 +63,8 @@ def __init__(  # type: ignore[no-untyped-def]
     @rpc
     def start(self) -> None:
         super().start()
-        self._disposables.add(self.color_image.subscribe(self._on_image))  # type: ignore[arg-type]
-        self._disposables.add(self.answer_stream.subscribe(self._on_answer))  # type: ignore[arg-type]
+        self.register_disposable(Disposable(self.color_image.subscribe(self._on_image)))
+        self.register_disposable(Disposable(self.answer_stream.subscribe(self._on_answer)))
         self._worker = threading.Thread(target=self._run_queries, daemon=True)
         self._worker.start()
 

dimos/agents/web_human_input.py

@@ -64,11 +64,11 @@ def start(self) -> None:
         # Subscribe to both text input sources
         # 1. Direct text from web interface
         unsub = self._web_interface.query_stream.subscribe(self._human_transport.publish)
-        self._disposables.add(unsub)
+        self.register_disposable(unsub)
 
         # 2. Transcribed text from STT
         unsub = stt_node.emit_text().subscribe(self._human_transport.publish)
-        self._disposables.add(unsub)
+        self.register_disposable(unsub)
 
         self._thread = Thread(target=self._web_interface.run, daemon=True)
         self._thread.start()

dimos/core/module.py

@@ -30,13 +30,12 @@
 )
 
 from langchain_core.tools import tool
-from reactivex.disposable import CompositeDisposable
 
 from dimos.core.core import T, rpc
 from dimos.core.global_config import GlobalConfig, global_config
 from dimos.core.introspection.module.info import extract_module_info
 from dimos.core.introspection.module.render import render_module_io
-from dimos.core.resource import Resource
+from dimos.core.resource import CompositeResource
 from dimos.core.rpc_client import RpcCall
 from dimos.core.stream import In, Out, RemoteOut, Transport
 from dimos.protocol.rpc.pubsubrpc import LCMRPC
@@ -92,15 +91,14 @@ class _BlueprintPartial(Protocol):
     def __call__(self, **kwargs: Any) -> "Blueprint": ...
 
 
-class ModuleBase(Configurable[ModuleConfigT], Resource):
+class ModuleBase(Configurable[ModuleConfigT], CompositeResource):
     # This won't type check against the TypeVar, but we need it as the default.
     default_config: type[ModuleConfigT] = ModuleConfig  # type: ignore[assignment]
 
     _rpc: RPCSpec | None = None
     _tf: TFSpec[Any] | None = None
     _loop: asyncio.AbstractEventLoop | None = None
     _loop_thread: threading.Thread | None
-    _disposables: CompositeDisposable
     _bound_rpc_calls: dict[str, RpcCall] = {}
     _module_closed: bool = False
     _module_closed_lock: threading.Lock
@@ -111,7 +109,6 @@ def __init__(self, config_args: dict[str, Any]):
         super().__init__(**config_args)
         self._module_closed_lock = threading.Lock()
         self._loop, self._loop_thread = get_loop()
-        self._disposables = CompositeDisposable()
         try:
             self.rpc = self.config.rpc_transport()
             self.rpc.serve_module_rpc(self)
@@ -132,6 +129,7 @@ def start(self) -> None:
 
     @rpc
     def stop(self) -> None:
+        super().stop()
         self._close_module()
 
     def _close_module(self) -> None:
@@ -158,14 +156,12 @@ def _close_module(self) -> None:
         if hasattr(self, "_tf") and self._tf is not None:
             self._tf.stop()
             self._tf = None
-        if hasattr(self, "_disposables"):
-            self._disposables.dispose()
 
-        # Break the In/Out -> owner -> self reference cycle so the instance
-        # can be freed by refcount instead of waiting for GC.
-        for attr in list(vars(self).values()):
-            if isinstance(attr, (In, Out)):
-                attr.owner = None
+        # Stop transports and break the In/Out -> owner -> self reference
+        # cycle so the instance can be freed by refcount instead of waiting for GC.
+        for attr in [*self.inputs.values(), *self.outputs.values()]:
+            attr.stop()
+            attr.owner = None
 
     def _close_rpc(self) -> None:
         if self.rpc:
@@ -188,7 +184,6 @@ def __setstate__(self, state) -> None:  # type: ignore[no-untyped-def]
         """Restore object from pickled state."""
         self.__dict__.update(state)
         # Reinitialize runtime attributes
-        self._disposables = CompositeDisposable()
         self._module_closed_lock = threading.Lock()
         self._loop = None
         self._loop_thread = None

dimos/core/resource.py

@@ -16,7 +16,7 @@
 
 from abc import abstractmethod
 import sys
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeVar
 
 if sys.version_info >= (3, 11):
     from typing import Self
@@ -29,6 +29,8 @@
 from reactivex.abc import DisposableBase
 from reactivex.disposable import CompositeDisposable
 
+D = TypeVar("D", bound=DisposableBase)
+
 
 class Resource(DisposableBase):
     @abstractmethod
@@ -75,18 +77,17 @@ def __exit__(
 class CompositeResource(Resource):
     """Resource that owns child disposables, disposed on stop()."""
 
-    _disposables: CompositeDisposable
-
-    def __init__(self) -> None:
-        self._disposables = CompositeDisposable()
+    _disposables: CompositeDisposable | None = None
 
-    def register_disposables(self, *disposables: DisposableBase) -> None:
-        """Register child disposables to be disposed when this resource stops."""
-        for d in disposables:
-            self._disposables.add(d)
+    def register_disposable(self, disposable: D) -> D:
+        """Register a child disposable to be disposed when this resource stops."""
+        if self._disposables is None:
+            self._disposables = CompositeDisposable()
+        self._disposables.add(disposable)
+        return disposable
 
-    def start(self) -> None:
-        pass
+    def start(self) -> None: ...
 
     def stop(self) -> None:
-        self._disposables.dispose()
+        if self._disposables is not None:
+            self._disposables.dispose()

dimos/core/stream.py

@@ -135,6 +135,10 @@ def __str__(self) -> str:
             + ("" if not self._transport else " via " + str(self._transport))
         )
 
+    def stop(self) -> None:
+        if self._transport is not None:
+            self._transport.stop()
+
 
 class Out(Stream[T], ObservableMixin[T]):
     _transport: Transport  # type: ignore[type-arg]

dimos/core/test_core.py

@@ -47,7 +47,7 @@ def _odom(msg) -> None:
             self.mov.publish(msg.position)
 
         unsub = self.odometry.subscribe(_odom)
-        self._disposables.add(Disposable(unsub))
+        self.register_disposable(Disposable(unsub))
 
         def _lidar(msg) -> None:
             self.lidar_msg_count += 1
@@ -57,7 +57,7 @@ def _lidar(msg) -> None:
                 print("RCV: unknown time", msg)
 
         unsub = self.lidar.subscribe(_lidar)
-        self._disposables.add(Disposable(unsub))
+        self.register_disposable(Disposable(unsub))
 
 
 def test_classmethods() -> None:

dimos/hardware/sensors/camera/module.py

@@ -77,11 +77,11 @@ def on_image(image: Image) -> None:
             self.color_image.publish(image)
             self._latest_image = image
 
-        self._disposables.add(
+        self.register_disposable(
             stream.subscribe(on_image),
         )
 
-        self._disposables.add(
+        self.register_disposable(
             rx.interval(1.0).subscribe(lambda _: self.publish_metadata()),
         )
 

dimos/hardware/sensors/camera/realsense/camera.py

@@ -162,15 +162,15 @@ def start(self) -> None:
 
         if self.config.enable_pointcloud and self.config.enable_depth:
             interval_sec = 1.0 / self.config.pointcloud_fps
-            self._disposables.add(
+            self.register_disposable(
                 backpressure(rx.interval(interval_sec)).subscribe(
                     on_next=lambda _: self._generate_pointcloud(),
                     on_error=lambda e: print(f"Pointcloud error: {e}"),
                 )
             )
 
         interval_sec = 1.0 / self.config.camera_info_fps
-        self._disposables.add(
+        self.register_disposable(
             rx.interval(interval_sec).subscribe(
                 on_next=lambda _: self._publish_camera_info(),
                 on_error=lambda e: print(f"CameraInfo error: {e}"),

dimos/hardware/sensors/camera/zed/camera.py

@@ -180,7 +180,7 @@ def start(self) -> None:
             self._enable_tracking()
 
         interval_sec = 1.0 / self.config.camera_info_fps
-        self._disposables.add(
+        self.register_disposable(
             rx.interval(interval_sec).subscribe(
                 on_next=lambda _: self._publish_camera_info(),
                 on_error=lambda e: print(f"CameraInfo error: {e}"),
@@ -193,7 +193,7 @@ def start(self) -> None:
 
         if self.config.enable_pointcloud and self.config.enable_depth:
             interval_sec = 1.0 / self.config.pointcloud_fps
-            self._disposables.add(
+            self.register_disposable(
                 backpressure(rx.interval(interval_sec)).subscribe(
                     on_next=lambda _: self._generate_pointcloud(),
                     on_error=lambda e: print(f"Pointcloud error: {e}"),